Inserting and Storing Embeddings in ChromaDB

Welcome back! In the previous lesson, you learned how to set up and initialize ChromaDB, a lightweight open-source vector database. You also created a collection to manage your vector data. In this lesson, we will build on that foundation by focusing on embeddings, which are crucial for converting text into numerical representations that can be efficiently stored and queried in ChromaDB. Our goal is to guide you through the process of inserting and storing these embeddings in ChromaDB, a key step in managing vector data for applications like semantic search.

To work with embeddings, we first need to load a pre-trained Sentence Transformer model. This model will help us convert text into vector representations. For this lesson, we will use the "sentence-transformers/all-MiniLM-L6-v2" model, which is known for its efficiency and accuracy in generating embeddings. You can load this model using the SentenceTransformer class from the sentence_transformers library. Here's how you can do it:

from sentence_transformers import SentenceTransformer

# Load embedding model
model_name = "sentence-transformers/all-MiniLM-L6-v2"
model = SentenceTransformer(model_name)

In this code snippet, we import the SentenceTransformer class and specify the model name. The model variable now holds the loaded model, ready to generate embeddings from text.

With the model loaded, the next step is to create a collection in ChromaDB that utilizes an embedding function. This function will transform text into vector representations before storing them in the database. We use the embedding_functions module from chromadb.utils to create an embedding function that leverages our loaded model. Here's how you can create a collection with an embedding function:

from chromadb.utils import embedding_functions

# Create embedding function
embed_func = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)

# Create collection with embedding function
collection = client.get_or_create_collection(name="vector_collection", embedding_function=embed_func)

In this example, we define an embedding function using the SentenceTransformerEmbeddingFunction class, passing the model name as a parameter. We then create or load a collection named "vector_collection" with this embedding function. This setup ensures that any text inserted into the collection is automatically converted into embeddings.