How to Use Cohere's Quantized Vectors to Build Cost-effective AI Apps With MongoDB

1
    Min: 1.23
2
    Max: 6.78

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

1
    quantized = round((original - Min) / Scale)

1
    def quantize_vector(vector, bits=8):
2
      """Quantize a float32 vector to n-bit integers."""
3
4
      # Convert the vector to a numpy array
5
      vector = np.array(vector, dtype=np.float32)
6
7
      # Step 1: Determine the range of values.
8
      min_val, max_val = vector.min(), vector.max()
9
10
      # Step 2: Define the quantization scale
11
      scale = (max_val - min_val) / (2**bits - 1)
12
      
13
      # Step 3: Quantize to 8-bit integers using the formula below: quantized = round((original - Min) / Scale)
14
      quantized = np.round((vector - min_val) / scale).astype(np.uint8)
15
      
16
      return quantized, min_val, scale

1
    def dequantize_vector(quantized, min_val, scale, bits=8):
2
      """Dequantize the vector back to float32."""
3
      return (quantized.astype(np.float32) * scale) + min_val
4
5
    # Original float32 vector
6
    original = np.array([1.23, 4.56, 2.34, 5.67, 3.45, 6.78, 1.89, 4.90], dtype=np.float32)
7
8
    # Run Quantize
9
    quantized, min_val, scale = quantize_vector(original)
10
11
    # Dequantize (reconstruct)
12
    reconstructed = dequantize_vector(quantized, min_val, scale)
13
14
    # Print results
15
    print("Original (float32):")
16
    print(original)
17
    print("\nQuantized (uint8):")
18
    print(quantized)
19
    print(f"\nQuantization parameters - Min: {min_val}, Scale: {scale}")
20
    print("\nReconstructed (float32):")
21
    print(reconstructed)
22
    print("\nMean Absolute Error:")
23
    print(np.mean(np.abs(original - reconstructed)))
24
25
    # Memory usage
26
    print(f"\nOriginal memory usage: {original.nbytes} bytes")
27
    print(f"Quantized memory usage: {quantized.nbytes} bytes")
28
    print(f"Compression ratio: {original.nbytes / quantized.nbytes:.2f}x")

1
    !pip install --quiet pandas cohere pymongo

1
    import os
2
    import cohere
3
    import getpass
4
5
    COHERE_API_KEY = getpass.getpass("Enter Cohere API Key: ")
6
    os.environ["COHERE_API_KEY"] = COHERE_API_KEY

1
    sentences = [
2
      {
3
          "sentence": "Water boils at 100 degrees Celsius at standard atmospheric pressure.",
4
          "float32_embedding": None,
5
          "int8_embedding": None,
6
          "bson_float32_embedding": None,
7
          "bson_int8_embedding": None
8
      },
9
      {
10
          "sentence": "The Great Wall of China is visible from space.",
11
          "float32_embedding": None,
12
          "int8_embedding": None,
13
          "bson_float32_embedding": None,
14
          "bson_int8_embedding": None
15
      },
16
      {
17
          "sentence": "Photosynthesis converts light energy into chemical energy in plants.",
18
          "float32_embedding": None,
19
          "int8_embedding": None,
20
          "bson_float32_embedding": None,
21
          "bson_int8_embedding": None
22
      },
23
      {
24
          "sentence": "DNA contains the genetic instructions for all living organisms.",
25
          "float32_embedding": None,
26
          "int8_embedding": None,
27
          "bson_float32_embedding": None,
28
          "bson_int8_embedding": None
29
      },
30
      {
31
          "sentence": "Jupiter is the largest planet in our solar system.",
32
          "float32_embedding": None,
33
          "int8_embedding": None,
34
          "bson_float32_embedding": None,
35
          "bson_int8_embedding": None
36
      },
37
      {
38
          "sentence": "The Eiffel Tower was completed in Paris in 1889.",
39
          "float32_embedding": None,
40
          "int8_embedding": None,
41
          "bson_float32_embedding": None,
42
          "bson_int8_embedding": None
43
      },
44
      {
45
          "sentence": "Honey never spoils due to its unique chemical properties.",
46
          "float32_embedding": None,
47
          "int8_embedding": None,
48
          "bson_float32_embedding": None,
49
          "bson_int8_embedding": None
50
      },
51
      {
52
          "sentence": "The human body contains approximately 206 bones in adulthood.",
53
          "float32_embedding": None,
54
          "int8_embedding": None,
55
          "bson_float32_embedding": None,
56
          "bson_int8_embedding": None
57
      },
58
      {
59
          "sentence": "Mount Everest is the highest peak on Earth at 8,848m.",
60
          "float32_embedding": None,
61
          "int8_embedding": None,
62
          "bson_float32_embedding": None,
63
          "bson_int8_embedding": None
64
      },
65
      {
66
          "sentence": "Shakespeare wrote 37 plays and 154 sonnets during his lifetime.",
67
          "float32_embedding": None,
68
          "int8_embedding": None,
69
          "bson_float32_embedding": None,
70
          "bson_int8_embedding": None
71
      },
72
      {
73
          "sentence": "The speed of light in vacuum is 299,792,458 meters/second.",
74
          "float32_embedding": None,
75
          "int8_embedding": None,
76
          "bson_float32_embedding": None,
77
          "bson_int8_embedding": None
78
      },
79
      {
80
          "sentence": "Penguins are flightless birds found primarily in the Southern Hemisphere.",
81
          "float32_embedding": None,
82
          "int8_embedding": None,
83
          "bson_float32_embedding": None,
84
          "bson_int8_embedding": None
85
      },
86
      {
87
          "sentence": "The Mona Lisa was painted by Leonardo da Vinci.",
88
          "float32_embedding": None,
89
          "int8_embedding": None,
90
          "bson_float32_embedding": None,
91
          "bson_int8_embedding": None
92
      },
93
      {
94
          "sentence": "Oxygen makes up about 21% of Earth's atmosphere by volume.",
95
          "float32_embedding": None,
96
          "int8_embedding": None,
97
          "bson_float32_embedding": None,
98
          "bson_int8_embedding": None
99
      },
100
      {
101
          "sentence": "The American Civil War lasted from 1861 to 1865.",
102
          "float32_embedding": None,
103
          "int8_embedding": None,
104
          "bson_float32_embedding": None,
105
          "bson_int8_embedding": None
106
      },
107
      {
108
          "sentence": "Antibiotics are ineffective against viral infections like the common cold.",
109
          "float32_embedding": None,
110
          "int8_embedding": None,
111
          "bson_float32_embedding": None,
112
          "bson_int8_embedding": None
113
      },
114
      {
115
          "sentence": "The human heart beats approximately 100,000 times per day.",
116
          "float32_embedding": None,
117
          "int8_embedding": None,
118
          "bson_float32_embedding": None,
119
          "bson_int8_embedding": None
120
      },
121
      {
122
          "sentence": "Gold is one of the least reactive chemical elements.",
123
          "float32_embedding": None,
124
          "int8_embedding": None,
125
          "bson_float32_embedding": None,
126
          "bson_int8_embedding": None
127
      },
128
      {
129
          "sentence": "The first successful powered aircraft flew in 1903.",
130
          "float32_embedding": None,
131
          "int8_embedding": None,
132
          "bson_float32_embedding": None,
133
          "bson_int8_embedding": None
134
      },
135
      {
136
          "sentence": "There are 118 elements in the periodic table.",
137
          "float32_embedding": None,
138
          "int8_embedding": None,
139
          "bson_float32_embedding": None,
140
          "bson_int8_embedding": None
141
      }
142
    ]

1
    import pandas as pd
2
3
    # Convert the list of dictionaries into a pandas DataFrame
4
    sentences_df = pd.DataFrame(sentences)
5
6
    sentences_df.head()

1
    import cohere
2
    from typing import List, Tuple
3
4
    # Initialize Cohere Client
5
    co = cohere.Client(COHERE_API_KEY)
6
7
    def get_cohere_embeddings(sentences: List[str]) -> Tuple[List[float], List[int]]:
8
        """
9
        Generates embeddings for the provided sentences using Cohere's embedding model.
10
11
        Args:
12
            sentences (list of str): List of sentences to generate embeddings for.
13
14
        Returns:
15
            Tuple[List[float], List[int]]: A tuple containing two lists of embeddings (float and int8).
16
        """
17
        generated_embedding = co.embed(
18
            texts=sentences,
19
            model="embed-english-v3.0",
20
            input_type="search_document",
21
            embedding_types=["float", "int8"]
22
        ).embeddings
23
24
        return generated_embedding.float, generated_embedding.int8

1
    # Iterate through each row in the DataFrame
2
    for index, row in sentences_df.iterrows():
3
        # Fetch the embeddings for the sentence
4
        embeddings = get_cohere_embeddings([row["sentence"]])
5
6
        # Assign the embeddings to the corresponding columns in the DataFrame
7
        sentences_df.at[index, "float32_embedding"] = embeddings[0][0]  # float32_embedding
8
        sentences_df.at[index, "int8_embedding"] = embeddings[1][0]  # int8_embedding

1
    sentences_df.head()

1
    from bson.binary import Binary, BinaryVectorDtype
2
    import numpy as np
3
    def generate_bson_vector(array, data_type):
4
        return Binary.from_vector(array, BinaryVectorDtype(data_type))
5
6
    # For all data points in the sentences_df generate BSON vectors for float32 and int8 embeddings
7
    sentences_df["bson_float32_embedding"] = sentences_df["float32_embedding"].apply(lambda x: generate_bson_vector(np.array(x, dtype=np.float32), BinaryVectorDtype.FLOAT32))
8
    sentences_df["bson_int8_embedding"] = sentences_df["int8_embedding"].apply(lambda x: generate_bson_vector(np.array(x, dtype=np.int8), BinaryVectorDtype.INT8))

1
    # remove the coloumn int8_embedding from the dataframe
2
    sentences_df = sentences_df.drop(columns=["int8_embedding"])

1
    sentences_df.head()

1
    MONGO_URI = getpass.getpass("Enter MongoDB URI: ")
2
    os.environ["MONGO_URI"] = MONGO_URI

1
    import pymongo
2
3
    def get_mongo_client(mongo_uri):
4
      """Establish and validate connection to the MongoDB."""
5
6
      client = pymongo.MongoClient(mongo_uri, appname="devrel.showcase.quantized_cohere.python")
7
8
      # Validate the connection
9
      ping_result = client.admin.command('ping')
10
      if ping_result.get('ok') == 1.0:
11
        # Connection successful
12
        print("Connection to MongoDB successful")
13
        return client
14
      else:
15
        print("Connection to MongoDB failed")
16
      return None
17
18
    if not MONGO_URI:
19
      print("MONGO_URI not set in environment variables")

1
    mongo_client = get_mongo_client(MONGO_URI)
2
3
    db = mongo_client["quantized"]
4
    col = db["data"]

1
    # Delete existing collection
2
    col.delete_many({})

1
    # Programmatically create vector search index for both colelctions
2
    import time
3
    from pymongo.operations import SearchIndexModel
4
5
    def setup_vector_search_index(collection, index_definition, index_name="vector_index"):
6
      """
7
      Setup a vector search index for a MongoDB collection and wait for 30 seconds.
8
9
      Args:
10
      collection: MongoDB collection object
11
      index_definition: Dictionary containing the index definition
12
      index_name: Name of the index (default: "vector_index")
13
      """
14
      new_vector_search_index_model = SearchIndexModel(
15
          definition=index_definition,
16
          name=index_name,
17
          type="vectorSearch"
18
      )
19
20
      # Create the new index
21
      try:
22
        result = collection.create_search_index(model=new_vector_search_index_model)
23
        print(f"Creating index '{index_name}'...")
24
25
        # Sleep for 30 seconds
26
        print(f"Waiting for 30 seconds to allow index '{index_name}' to be created...")
27
        time.sleep(30)
28
29
        print(f"30-second wait completed for index '{index_name}'.")
30
        return result
31
32
      except Exception as e:
33
        print(f"Error creating new vector search index '{index_name}': {str(e)}")
34
        return None

1
    dimensions = 1024 # Output dimensions for vectors generated by Cohere's embed-english-v3.0 embedding model
2
3
    vector_search_index_definition = {
4
      "fields": [
5
        {
6
          "type": "vector",
7
          "path": "float32_embedding",
8
          "numDimensions": dimensions,
9
          "similarity": "cosine"
10
        },
11
        {
12
          "type": "vector",
13
          "path": "bson_float32_embedding",
14
          "numDimensions": dimensions,
15
          "similarity": "cosine"
16
        },
17
        {
18
          "type": "vector",
19
          "path": "bson_int8_embedding",
20
          "numDimensions": dimensions,
21
          "similarity": "cosine"
22
        }
23
      ]
24
    }

1
    setup_vector_search_index(col, vector_search_index_definition, index_name="vector_index")

1
    Creating index 'vector_index'...
2
    Waiting for 30 seconds to allow index 'vector_index' to be created...
3
    30-second wait completed for index 'vector_index'.
4
    vector_index

1
    def vector_search(user_query, collection, quantized=False, use_bson=True):
2
      """
3
      Perform a vector search in the MongoDB collection based on the user query.
4
5
      Args:
6
      user_query (str): The user's query string.
7
      collection (MongoCollection): The MongoDB collection to search.
8
9
      Returns:
10
      list: A list of matching documents.
11
      """
12
13
      path = None
14
      query_embedding = None
15
16
      # Generate embedding for the user query
17
      query_embedding = get_cohere_embeddings([user_query])
18
19
      if query_embedding is None:
20
          return "Invalid query or embedding generation failed."
21
22
      if quantized:
23
        query_embedding = query_embedding[1][0]
24
        path = "int8_embedding"
25
      else:
26
        query_embedding = query_embedding[0][0]
27
        path = "float32_embedding"
28
29
30
       # Adjust path and convert vectors to BSON vectors if use_bson is True
31
      if use_bson:
32
        path = "bson_" + path
33
        if quantized:
34
          query_embedding = generate_bson_vector(query_embedding, BinaryVectorDtype.INT8)
35
        else:
36
          query_embedding = generate_bson_vector(query_embedding, BinaryVectorDtype.FLOAT32)
37
38
39
      # Define the vector search pipeline
40
      vector_search_stage = {
41
        "$vectorSearch": {
42
          "index": "vector_index",
43
          "queryVector": query_embedding,
44
          "path": path,
45
          "numCandidates": 150,  # Number of candidate matches to consider
46
          "limit": 5  # Return top 5 matches
47
        }
48
      }
49
50
      project_stage = {
51
        "$project": {
52
          "_id": 0,  # Exclude the _id field
53
          "sentence": 1,  # Include the sentence field
54
          "similarityScore": {
55
            "$meta": "vectorSearchScore"  # Include the search score
56
          }
57
        }
58
      }
59
60
      pipeline = [vector_search_stage, project_stage]
61
62
      # Execute the search
63
      results = collection.aggregate(pipeline)
64
      return list(results)

1
    results_bson_float = vector_search("what is the speed of light?", col, quantized=False, use_bson=True)
2
    results_bson_int8 = vector_search("what is the speed of light?", col, quantized=True, use_bson=True)
3
    results_float = vector_search("what is the speed of light?", col, quantized=False, use_bson=False)

1
    import pprint
2
3
    print("Float32 Results")
4
    pprint.pprint(results_float)
5
    print("-----------------")
6
    print()
7
    print("BSON Float32 Results")
8
    pprint.pprint(results_bson_float)
9
    print("-----------------")
10
    print()
11
    print("BSON Int8 Results")
12
    pprint.pprint(results_bson_int8)
13
    print("-----------------")