Saving and Using Embeddings Locally in Python

In this lesson, we will explore the process of saving and using embeddings locally using Python. By focusing on practical implementation, you will gain hands-on experience in managing embeddings within your projects. By the end of this lesson, you will understand how to save embeddings for future use and load them back into your application, allowing you to streamline your workflows and improve the efficiency of your systems.

Saving embeddings is crucial as it allows you to reuse them without recalculating, thereby conserving computational resources and time. This is especially important when dealing with large datasets or when embeddings are generated using models that require significant processing power. In production environments and with very large databases, embeddings are typically stored in specialized databases known as vector databases. However, in this lesson, we will focus on saving embeddings locally, which is a practical approach for smaller projects or initial development stages. We will explore different vector databases in future courses.

First, let's define our sentences and calculate their embeddings using a pre-trained model.

from sentence_transformers import SentenceTransformer

# Define sentences
sentences = ["This is a sample sentence.", "Embeddings are useful for NLP tasks."]

# Load a pre-trained model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Calculate embeddings
embeddings = model.encode(sentences)

After generating the embeddings, it's important to save them for future use. In this section, we'll demonstrate how to store embeddings in both JSON and Pickle formats. JSON is a widely-used, human-readable format that makes it easy to inspect and share data. In contrast, Pickle is a binary format that is optimized for quick storage and retrieval, making it particularly useful when working with large datasets.

The format below is a minimal demo. In production, you would typically store metadata such as a stable text id, model_name, embedding_dim, creation timestamp, and the mapping between each sentence and its embedding.

```python

```python
import json
import pickle

# Save embeddings to JSON
# We'll store each sentence and its embedding in a list of dicts
data_to_save = []
for text, emb in zip(sentences, embeddings):
    data_to_save.append({
        "sentence": text,
        "embedding": emb.tolist()  # Convert NumPy array to Python list
    })

with open('embeddings.json', 'w', encoding='utf-8') as f:
    json.dump(data_to_save, f, indent=2)  # Pretty-print with indent

# Save embeddings to Pickle
with open('embeddings.pkl', 'wb') as f:
    pickle.dump(embeddings, f)

In this code, we first convert the embeddings into a list of dictionaries, where each embedding is converted from a NumPy array to a Python list to ensure compatibility with JSON serialization. This allows for straightforward inspection and sharing. We also save the embeddings in Pickle format, which is efficient for both storage and retrieval.

This is a minimal demo format. In a production setting, save metadata such as a stable text ID, source sentence or document ID, model name, embedding dimension, and creation timestamp so embeddings can be traced back to their source and model version.

To utilize the saved embeddings, we need to load them back into our application. Below are the steps to load embeddings from both JSON and Pickle files.

First, let's load the embeddings stored in the JSON file:

# Load embeddings from JSON
with open('embeddings.json', 'r', encoding='utf-8') as f:
    loaded_data = json.load(f)
    loaded_embeddings_json = [item['embedding'] for item in loaded_data]

print(f"Loaded Embedding from JSON (first sentence): {loaded_embeddings_json[0][:5]}...")

# Output:
# Loaded Embedding from JSON (first sentence): [0.07780639827251434, 0.07646248489618301, 0.0377088338136673, 0.06093894690275192, 0.04880760237574577]...

In this code snippet, we extract the embeddings from the loaded data, which can then be used for various tasks.

Next, let's load the embeddings from the Pickle file:

# Load embeddings from Pickle
with open('embeddings.pkl', 'rb') as f:
    loaded_embeddings_pickle = pickle.load(f)

print(f"Loaded Embedding from Pickle (first sentence): {loaded_embeddings_pickle[0][:5]}...")

# Output:
# Loaded Embedding from Pickle (first sentence): [0.0778064  0.07646248 0.03770883 0.06093895 0.0488076 ]...

This code demonstrates how to load embeddings from a Pickle file into a variable.

Both methods enable you to reuse the embeddings for tasks such as calculating similarities or integrating them into machine learning models.

For larger systems, vector databases provide optimized storage and querying for embeddings instead of loading local files into memory. We will cover Qdrant in a later course when learners build applications that actually store and retrieve embeddings at scale.

In this lesson, we focused on the process of saving and using embeddings locally using Python. We explored how to store embeddings in both JSON and Pickle formats, allowing for efficient reuse in various applications. Additionally, we introduced the concept of using vector databases like Qdrant for handling large-scale applications, which provides optimized storage and querying solutions for high-dimensional vector data. By storing embeddings, you can save time and computational resources, which is crucial for building scalable and efficient systems. As you move forward, consider practicing by saving embeddings for your own datasets and experimenting with different storage formats to find what best suits your needs. Good job on completing this lesson—you're almost there! Let's now tackle the final practice tasks.