Efficient Quantized Vector Ingestion with Cohere and MongoDB

You can view an article version of this notebook here:

In this notebook, we cover the following:

• What are quantization and vector quantization?
• Example of scalar quantization
• Efficient quantized vector ingestion with Cohere and MongoDB
• Comparison of float32, BSON float32, and BSON int8 embeddings in vector search operations
• Practical considerations for choosing the right embedding format for your AI application

Prerequisute: Understanding Scalar Quantization

Vector quantization is a lossy compression algorithm designed to reduce the dimensionality and memory requirements of high-dimensional vector data. It achieves this by mapping the original vectors to a smaller set of representative vectors. This process allows for significant data compression while preserving much of the essential information contained in the original vectors. Vector quantization technique is particularly useful in AI applications where some loss of precision is acceptable in exchange for reduced storage needs and lower latency in data retrieval steps. Let's demonstrate vector quantization using float32 values and quantizing to 8-bit integers. We'll use a simplified example with 1D vectors for clarity.

	Float32: [1.23, 4.56, 2.34, 5.67, 3.45, 6.78, 1.89, 4.90]

	Int8: [0, 153, 51, 204, 102, 255, 30, 169]

Step 1: Determine the range of values.

	Min: 1.23
Max: 6.78

Step 2: Define the quantization scale

	Scale = (Max - Min) / 255 = (6.78 - 1.23) / 255 ≈ 0.0217

Step 3: Quantize to 8-bit integers using the formula below: quantized = round((original - Min) / Scale)

What we are building

This section provides a comprehensive guide to implementing efficient quantized vector ingestion and search using Cohere and MongoDB. We'll walk through the entire process, from generating quantized and non-quantized embeddings with Cohere's API to converting them into BSON format, ingesting them into MongoDB, and performing vector searches on the stored embeddings.

The end result of this section is a fully functional system that demonstrates the practical implementation of quantized vector storage and retrieval. This setup will allow you to compare the performance and accuracy of different embedding types (float32, BSON float32, and BSON int8) in vector search operations.

Step 1: Install Libraries and Set Environment Variables

For the code implementation in this tutorial, the following libraries are utilized:

• pandas:  A data manipulation library for efficient handling of structured data. It's used for loading, cleaning, transforming, and analyzing data in various formats.
• cohere: The official Cohere Python library. It will provide access to advanced language models, embedding generation, and text generation.
• pymongo: The official Python driver for MongoDB. While commented out in the installation, it suggests potential use for interacting with MongoDB databases, enabling storage and retrieval of data.

Click here to get your Cohere API Key

Step 2: Data Loading

This step involves generating a small structured dataset of sentences with placeholder fields for various embedding types. The dataset is designed to facilitate the demonstration and comparison of different embedding techniques in natural language processing and database storage.

The data is structured as a list of dictionaries, where each dictionary represents a sentence and its associated embedding attributes. Each entry contains a sentence key with a factual statement as its value, along with four embedding-related keys: float32_embedding, int8_embedding, bson_float32_embedding, and bson_int8_embedding. These keys are initially set to None.

Step 3: Embedding Generation With Cohere

This step demonstrates the process of generating embeddings for our sentence dataset using the Cohere API. We will define a custom function, get_cohere_embeddings, for efficient embedding generation.

Step 4: Generate BSON Vectors

This step introduces a crucial data transformation process that, while potentially unfamiliar to some developers, is integral to building high-performant AI applications. When optimizing AI systems for scale and efficiency, certain techniques and data formats become essential to application development. The operations demonstrated here are fundamental to showcasing the memory optimization benefits of BSON (Binary JSON) and MongoDB's storage capabilities.

The operations demonstrated here are fundamental to showcasing the memory optimization benefits of BSON (Binary JSON) and MongoDB's storage capabilities. By converting our embeddings into BSON format, we can compare the storage efficiency of the BSON data format vs the non-BSON data format.

Step 5: Data Ingestion With MongoDB

For this step, you will need a MongoDB account to retrieve a connection string to a cluster.

Creating a database and collection within MongoDB is simple with MongoDB Atlas.

1. First, register for a MongoDB Atlas account. Existing users can sign into MongoDB Atlas.
2. Follow the instructions. Select Atlas UI as the procedure to deploy your first cluster.
3. Follow MongoDB’s steps to get the connection string from the Atlas UI. Securely store the URI within your development environment after setting up the database and obtaining the Atlas cluster connection URI.

NOTE: Don’t forget to add the IP of your host machine to the IP access list for your cluster.

Step 6: Create Vector Index

In this step, we will create a critical component that facilitates vector search operations, a vector index. The setup_vector_search_index function creates a vector index, a specialized structure that significantly enhances the performance of similarity searches across embedding vectors.

The code snippet above utilizes MongoDB's SearchIndexModel to define a vector search index. One operation to highlight is the 30-second wait period post-index creation. This ensures the index is fully built and optimized, preventing potential race conditions in subsequent query operations.

The next step is to initialize the variable vector_search_index_definition that is assigned to the actual vector index, which is configured for multiple vector fields, each corresponding to different embedding formats (float32, BSON float32, and BSON int8). By setting up this index, we enable efficient cosine similarity searches across our 1024-dimensional embedding space.

Step 7: Create Vector Search Function

In this step, we'll define the vector_search function that leverages MongoDB's powerful aggregation framework to perform efficient vector search operations. This function is the intermediary layer between user queries and our database of sentence embeddings, enabling semantic search or RAG capabilities in your AI applications.

The vector_search function accommodates different types of embeddings and search strategies. It takes a user's input query, typically a sentence from the application layer, and transforms it into a vector representation using our previously defined get_cohere_embeddings function. This embedding, stored in the query_embedding variable, becomes the basis for our similarity search.

A key feature of this function is its ability to handle both quantized and non-quantized embeddings. By incorporating a check for the embedding type, we can dynamically adjust our search parameters to utilize either the space-efficient quantized (int8) embeddings or the higher-precision non-quantized (float32) embeddings.

By leveraging MongoDB's aggregation pipeline, we'll construct a series of stages that process our data, perform the vector similarity search, and return the most relevant results.

The key operations of the code snippet above and the vector_search function are as follows:

1. Embedding generation: The get_cohere_embeddings function converts the user's query into a vector representation.
2. Embedding type selection: Based on the quantized flag, it chooses between float32 and int8 embeddings, adjusting the search path accordingly.
3. BSON conversion: If use_bson is True, the embedding is converted to BSON format using the generate_bson_vector function, optimizing MongoDB storage and retrieval.
4. Search pipeline construction: The function builds a MongoDB aggregation pipeline with two stages:

• $vectorSearch: This utilizes the pre-created vector index for efficient similarity search. This operator takes the search query embedding as input, the path to the location of embedding data located in the collection, the number of candidates to consider for similarity, and the number of documents to be retrieved by the vector stage. -$project: This shapes the output, including the original sentence and a similarity score.

Step 8: Run Vector Search Query

In this step, we will use the vector_search function across different embedding configurations. This comparison allows us to evaluate the performance and accuracy of various vector representations in semantic search tasks.

In the code snippet above, we execute three distinct vector searches using the query "what is the speed of light?" each targeting a different embedding format:

1. results_bson_float: Uses BSON-encoded float32 embeddings, combining the precision of floating-point numbers with the efficiency of BSON storage
2. results_bson_int8: Utilizes BSON-encoded quantized (int8) embeddings, offering a more compact representation that may trade some precision for improved storage and query efficiency
3. `results_float`: Uses standard float32 embeddings without BSON encoding, serving as a baseline for comparison

These results will provide insights into how different embedding representations affect semantic search outcomes, helping you choose the most suitable approach for your specific use cases; typical factors to consider are storage optimization, search speed, and result precision.

Conclusion

All code presented in the steps outlined in this tutorial are available on GitHub.

Based on these results, we can draw several important conclusions about the benefits of quantization and BSON encoding in vector search operations:

We will start off with an important consideration, which is accuracy preservation. The results demonstrate that MongoDB’s native support for Vectors in BSON and quantization maintain a high degree of accuracy in semantic search. The top result for all three methods is identical, correctly identifying the sentence about the speed of light, with very similar similarity scores (0.8345 for float32 and BSON float32, 0.8344 for BSON int8).

Another factor to consider is the consistency across formats. The order of results and the similarity scores are remarkably consistent across all three formats. This suggests that the use of BSON encoding and even quantization to int8 does not significantly compromise the quality of search results. Based on this, let’s recap the benefits of BSON and quantized vectors for AI applications.

1. Benefits of native vector BSON type:

• Storage efficiency: While not directly visible in the results from this tutorial, BSON encoding typically offers more compact storage compared to raw float32 values, potentially reducing database size. Follow the advanced notebook to get an overview of BSON's storage efficiency.
• Query performance: BSON's binary format can lead to faster serialization/deserialization, potentially improving query speeds in large-scale applications.

2. Advantages of quantization (int8):

• Space efficiency: Int8 representation requires significantly less storage space compared to float32, which can be crucial for large datasets.
• Potential performance gains: Reduced data size can lead to faster data transfer and potentially quicker computations, which is especially beneficial for resource-constrained environments.

3. Minimal precision loss: The slight variations in similarity scores for the BSON int8 results (e.g., 0.8344 vs 0.8345 for the top result) indicate that the precision loss from quantization is minimal and does not affect the overall ranking of results.

The results of this tutorial demonstrate that you can leverage the benefits of both BSON encoding and quantization without significant compromise to search accuracy.

In an internal experiment, we compared the storage efficiency of different representations for vector embeddings: floats and BSON vectors using float32 precision. The test dataset consisted of three million documents, each containing embeddings generated by the cohere-embed-multilingual-v3.0 model with 1024 dimensions, like this tutorial. The results were as follows:

This comparison demonstrates a significant reduction in storage requirements when using BSON vectors with float32 precision. The BSON representation required only about 34% of the storage space needed for the array of floats representation, resulting in approximately 66% space savings. This substantial difference in storage efficiency has important implications for large-scale vector databases.

BSON encoding offers improved storage and potential performance benefits, while further quantization to int8 from float32 embeddings, as in this tutorial, can provide additional space savings with negligible impact on result quality.

The choice between these formats (array of floats, BSON vectors, or quantized representations) should be based on specific application needs, balancing storage constraints, processing speed requirements, and the level of precision required for the particular use case.

A key takeaway is that MongoDB Atlas enhances your AI stack by providing robust indexing support for quantized embedding model outputs from providers such as Cohere. This capability significantly improves the scalability of AI workloads, particularly those involving vector search, allowing them to handle substantially higher volumes. The smaller memory footprint of quantized embeddings results in more efficient storage utilization and cost-effectiveness. These features make MongoDB Atlas an excellent choice for organizations looking to optimize their AI applications, particularly those involving large-scale vector search operations.

This tutorial covers an introductory approach to evaluating the benefits of quantized embeddings. For a more detailed and advanced exploration of the memory optimization benefits of quantized vectors and BSON data format, check out the GitHub notebook.