Inserting and Storing Embeddings in ChromaDB | CodeSignal Learn

Inserting and Storing Embeddings in ChromaDB | CodeSignal Learn

Introduction to 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.

Loading the Sentence Transformer Model

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: Python 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.

Creating a Collection with an Embedding Function

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: Python 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.

Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Inserting Documents into ChromaDB

Now that we have a collection with an embedding function, we can insert documents into ChromaDB. Each document will be transformed into an embedding and stored in the collection. Let's walk through the process using a few sample documents: # Sample documents documents = [{"id": "doc1", "content": "ChromaDB is an open-source vector database."}, {"id": "doc2", "content": "Vector databases help in efficient similarity search."}, {"id": "doc3", "content": "Embedding models convert text into numerical representations."}] # Insert documents into ChromaDB collection.add(documents=[doc["content"] for doc in documents], ids=[doc["id"] for doc in documents]) print("Inserted documents into ChromaDB.") In this code, we define a list of sample documents, each with a unique identifier and content. We then use the add method of the collection to insert these documents. The documents parameter takes a list of text content, while the ids parameter takes a list of corresponding document IDs. When you run this code, you should see the output: "Inserted documents into ChromaDB", confirming that the documents have been successfully stored.

Retrieving Documents and Embeddings from ChromaDB

After inserting documents into ChromaDB, you may want to retrieve them to verify their storage or use them in further operations. The collection.get() method allows you to fetch documents from the collection using their IDs. Additionally, you can use the include parameter to specify that you want to retrieve embeddings along with the document content. This is useful for confirming that your documents have been correctly stored and for accessing their content and embeddings when needed. Here's how you can use the collection.get() method with the include parameter: # Retrieve and display embedding for a specific document doc_id = "doc1" result = collection.get(ids=[doc_id], include=["embeddings"]) if result and result.get("embeddings") is not None and len(result["embeddings"]) > 0: embedding = result["embeddings"][0] print(f"Embedding for {doc_id}: {embedding}") print(f"Embedding size: {len(embedding)}") else: print(f"No embedding found for document ID: {doc_id}") In this example, we use the get method to retrieve the document with ID "doc1", including its embedding. The method returns a dictionary containing the document ID, content, and embedding. We then check if the embedding is available and print it along with its size to verify the retrieval. This approach allows you to access both the content and the embeddings of your documents in ChromaDB.

Python

# Sample documents
documents = [
    {"id": "doc1", "content": "ChromaDB is an open-source vector database."},
    {"id": "doc2", "content": "Vector databases help in efficient similarity search."},
    {"id": "doc3", "content": "Embedding models convert text into numerical representations."}
]

# Insert documents into ChromaDB
collection.add(
    documents=[doc["content"] for doc in documents],
    ids=[doc["id"] for doc in documents]
)

print("Inserted documents into ChromaDB.")

Python

# Retrieve and display embedding for a specific document
doc_id = "doc1"
result = collection.get(ids=[doc_id], include=["embeddings"])

if result and result.get("embeddings") is not None and len(result["embeddings"]) > 0:
    embedding = result["embeddings"][0]
    print(f"Embedding for {doc_id}: {embedding}")
    print(f"Embedding size: {len(embedding)}")
else:
    print(f"No embedding found for document ID: {doc_id}")