Reducing Query Latency with Precomputed Nearest Neighbors

Introduction to Query Latency in Vector Search

Welcome to the first lesson of our course on " Optimizing and Scaling ChromaDB for Vector Search ." In this lesson, we will explore the concept of query latency in vector search systems and how it impacts user experience. Query latency refers to the time it takes for a search query to return results. In vector search systems, reducing this latency is crucial for providing a seamless and efficient user experience. One effective method to achieve this is by precomputing nearest neighbors, which allows us to quickly retrieve relevant results without recalculating distances for every query. This lesson will guide you through the process of implementing precomputed nearest neighbors using ChromaDB .

Overview of ChromaDB and Embedding Functions

ChromaDB is a powerful tool for managing and querying vector data. It allows us to store, index, and retrieve high-dimensional vectors efficiently. In this lesson, we will use ChromaDB to manage our vector data and perform nearest neighbor searches. An essential component of this process is the use of embedding functions . Embedding functions transform text data into vector representations, which can then be used for similarity searches. In our example, we will use the sentence-transformers/all-MiniLM-L6-v2 model to generate embeddings. This model is known for its efficiency and accuracy in creating meaningful vector representations of text. Python from chromadb import Client from chromadb.utils import embedding_functions # Set up Chroma model_name = "sentence-transformers/all-MiniLM-L6-v2" embed_func = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name) client = Client() collection = client.get_or_create_collection("document_collection", embedding_function=embed_func)

Loading Data

To begin, we need to load our data from a JSON file containing a corpus of documents. Each document includes content, metadata such as title and category, and a unique identifier. Here is an example of how a document in the JSON file might look: JSON{ "id": 1, "title": "Revolutionizing Computing with AI", "content": "Artificial intelligence is transforming the way we approach complex problems in computing. Recent breakthroughs in machine learning have enabled faster data processing and smarter algorithms. The future of technology is expected to integrate AI into every facet of life.", "category": "Technology", "tags": ["AI", "machine learning", "computing", "innovation"], "date": "2025-02-01" } We will use our loader function load_documents defined in the data directory to read the documents from the JSON file: Pythonfrom data import load_documents # Load documents documents = load_documents("./data/corpus.json") print(f"Loaded {len(documents)} documents.")

Join the 1M+ learners on CodeSignal

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

Batch Insertion into ChromaDB

To efficiently insert data into ChromaDB, we can use a batch insertion method. This approach allows us to add documents in chunks, reducing the load on the system and improving performance. Below is the code for batch insertion:

Python

# Batch add function
def batch_add_to_chroma(collection, documents, batch_size=50):
    print("Starting batch insert into ChromaDB...")
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        print(f"Inserting batch {i} to {i + len(batch)}...")
        collection.add(
            documents=[doc["content"] for doc in batch],
            ids=[str(doc["id"]) for doc in batch],
            metadatas=[{
                "title": doc["title"],
                "category": doc.get("category", "unknown"),
                "tags": ",".join(doc.get("tags", [])),
                "date": doc.get("date", "unknown")
            } for doc in batch]
        )
    print("Finished inserting all batches.")

# Add data to Chroma
batch_add_to_chroma(collection, documents)

# Batch add function
def batch_add_to_chroma(collection, documents, batch_size=50):
    print("Starting batch insert into ChromaDB...")
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        print(f"Inserting batch {i} to {i + len(batch)}...")
        collection.add(
            documents=[doc["content"] for doc in batch],
            ids=[str(doc["id"]) for doc in batch],
            metadatas=[{
                "title": doc["title"],
                "category": doc.get("category", "unknown"),
                "tags": ",".join(doc.get("tags", [])),
                "date": doc.get("date", "unknown")
            } for doc in batch]
        )
    print("Finished inserting all batches.")

# Add data to Chroma
batch_add_to_chroma(collection, documents)

This code initializes a ChromaDB client and creates a collection named "document_collection." It then loads the corpus from a JSON file and inserts the documents into the collection in batches, along with their metadata.

Precomputing Nearest Neighbors

Once our data is in ChromaDB, we can precompute the nearest neighbors for each document. This involves calculating the cosine similarity between the vector representations of the documents. By precomputing these similarities, we can quickly retrieve the most similar documents for any given query, significantly reducing query latency. The following code demonstrates how to precompute nearest neighbors:

Python

from sklearn.metrics.pairwise import cosine_similarity

# Precompute nearest neighbors
precomputed_neighbors = {}

def precompute_neighbors(collection, embed_func, top_k=5):
    # Retrieve documents and their embeddings
    documents = collection.get()["documents"]
    embeddings = embed_func(documents)
    
    # Calculate cosine similarity matrix
    similarity_matrix = cosine_similarity(embeddings, embeddings)
    
    # Store top-k nearest neighbors for each document
    neighbors = {
        doc: sorted(
            zip(documents, similarity_matrix[i]),  # Pair documents with their similarity scores
            key=lambda x: x[1],  # Sort by similarity score
            reverse=True  # Highest similarity first
        )[:top_k]  # Keep only the top-k most similar documents
        for i, doc in enumerate(documents)  # Iterate over documents and their indices
    }
    return neighbors  

precomputed_neighbors = precompute_neighbors(collection, embed_func)
print("Precomputed nearest neighbors stored.")

from sklearn.metrics.pairwise import cosine_similarity

# Precompute nearest neighbors
precomputed_neighbors = {}

def precompute_neighbors(collection, embed_func, top_k=5):
    # Retrieve documents and their embeddings
    documents = collection.get()["documents"]
    embeddings = embed_func(documents)
    
    # Calculate cosine similarity matrix
    similarity_matrix = cosine_similarity(embeddings, embeddings)
    
    # Store top-k nearest neighbors for each document
    neighbors = {
        doc: sorted(
            zip(documents, similarity_matrix[i]),  # Pair documents with their similarity scores
            key=lambda x: x[1],  # Sort by similarity score
            reverse=True  # Highest similarity first
        )[:top_k]  # Keep only the top-k most similar documents
        for i, doc in enumerate(documents)  # Iterate over documents and their indices
    }
    return neighbors  

precomputed_neighbors = precompute_neighbors(collection, embed_func)
print("Precomputed nearest neighbors stored.")

In this code, we retrieve the documents and their embeddings from the collection, calculate the cosine similarity for each document against all others, and store the top 5 nearest neighbors. This precomputation allows for rapid retrieval of similar documents during search queries.

Let's walk through the complete example to see how all the pieces fit together. We start by loading our data into ChromaDB, then precompute the nearest neighbors using cosine similarity. This process involves initializing the ChromaDB client, loading the corpus, inserting documents, and finally, precomputing and storing the nearest neighbors. By following these steps, you can efficiently reduce query latency in your vector search system. The output of the precomputation step confirms that the nearest neighbors have been successfully stored, ready for quick retrieval during searches.

In this lesson, we explored the concept of query latency and how precomputing nearest neighbors can help reduce it in vector search systems. We learned how to use ChromaDB to manage vector data and perform efficient similarity searches. By precomputing nearest neighbors, we can significantly improve the performance of our search system. As you move on to the practice exercises, you will have the opportunity to reinforce these concepts and apply them to real-world scenarios. Remember, reducing query latency is crucial for providing a seamless user experience, and precomputing nearest neighbors is a powerful technique to achieve this. Good luck with your practice exercises!

Introduction to Query Latency in Vector Search

Overview of ChromaDB and Embedding Functions

Loading Data

Join the 1M+ learners on CodeSignal

Batch Insertion into ChromaDB

Precomputing Nearest Neighbors

Example: Implementing Precomputed Nearest Neighbors

Summary and Preparation for Practice Exercises