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.

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') 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.