Vector Quantization

pip install --quiet --upgrade pymongo cohere

data = [
   "The Great Wall of China is visible from space.",
   "The Eiffel Tower was completed in Paris in 1889.",
   "Mount Everest is the highest peak on Earth at 8,848m.",
   "Shakespeare wrote 37 plays and 154 sonnets during his lifetime.",
   "The Mona Lisa was painted by Leonardo da Vinci.",
]

import os
import cohere
# Specify your Cohere API key
os.environ["COHERE_API_KEY"] = "<COHERE-API-KEY>"
cohere_client = cohere.Client(os.environ["COHERE_API_KEY"])
# Generate embeddings using the embed-english-v3.0 model
generated_embeddings = cohere_client.embed(
   texts=data,
   model="embed-english-v3.0",
   input_type="search_document",
   embedding_types=["float", "int8", "ubinary"]
).embeddings
float32_embeddings = generated_embeddings.float
int8_embeddings = generated_embeddings.int8
int1_embeddings = generated_embeddings.ubinary

from bson.binary import Binary, BinaryVectorDtype
def generate_bson_vector(vector, vector_dtype):
   return Binary.from_vector(vector, vector_dtype)
# For all vectors in your collection, generate BSON vectors of float32, int8, and int1 embeddings
bson_float32_embeddings = []
bson_int8_embeddings = []
bson_int1_embeddings = []
for i, (f32_emb, int8_emb, int1_emb) in enumerate(zip(float32_embeddings, int8_embeddings, int1_embeddings)):
   bson_float32_embeddings.append(generate_bson_vector(f32_emb, BinaryVectorDtype.FLOAT32))
   bson_int8_embeddings.append(generate_bson_vector(int8_emb, BinaryVectorDtype.INT8))
   bson_int1_embeddings.append(generate_bson_vector(int1_emb, BinaryVectorDtype.PACKED_BIT))

# Specify the field names for the float32, int8, and int1 embeddings
float32_field = "<FIELD-NAME-FOR-FLOAT32-TYPE>"
int8_field = "<FIELD-NAME-FOR-INT8-TYPE>"
int1_field = "<FIELD-NAME-FOR-INT1-TYPE>"
# Define function to create documents with BSON vector embeddings
def create_docs_with_bson_vector_embeddings(bson_float32_embeddings, bson_int8_embeddings, bson_int1_embeddings, data):
  docs = []
  for i, (bson_f32_emb, bson_int8_emb, bson_int1_emb, text) in enumerate(zip(bson_float32_embeddings, bson_int8_embeddings, bson_int1_embeddings, data)):
     doc = {
          "_id": i,
          "data": text,
          float32_field: bson_f32_emb,
          int8_field: bson_int8_emb,
          int1_field: bson_int1_emb
     }
     docs.append(doc)
  return docs
# Create the documents
documents = create_docs_with_bson_vector_embeddings(bson_float32_embeddings, bson_int8_embeddings, bson_int1_embeddings, data)

from pymongo.operations import SearchIndexModel
import time
# Define and create the vector search index
index_name = "<INDEX-NAME>"
search_index_model = SearchIndexModel(
  definition={
    "fields": [
      {
        "type": "vector",
        "path": float32_field,
        "similarity": "dotProduct",
        "numDimensions": 1024
      },
      {
        "type": "vector",
        "path": int8_field,
        "similarity": "dotProduct",
        "numDimensions": 1024
      },
      {
        "type": "vector",
        "path": int1_field,
        "similarity": "euclidean",
        "numDimensions": 1024
      }
    ]
  },
  name=index_name,
  type="vectorSearch"
)
result = collection.create_search_index(model=search_index_model)
print("New search index named " + result + " is building.")
# Wait for initial sync to complete
print("Polling to check if the index is ready. This may take up to a minute.")
predicate=None
if predicate is None:
  predicate = lambda index: index.get("queryable") is True
while True:
  indices = list(collection.list_search_indexes(index_name))
  if len(indices) and predicate(indices[0]):
    break
  time.sleep(5)
print(result + " is ready for querying.")

# Define a function to run a vector search query
def run_vector_search(query_text, collection, path):
  query_text_embeddings = cohere_client.embed(
    texts=[query_text],
    model="embed-english-v3.0",
    input_type="search_query",
    embedding_types=["float", "int8", "ubinary"]
  ).embeddings
  if path == float32_field:
    query_vector = query_text_embeddings.float[0]
    vector_dtype = BinaryVectorDtype.FLOAT32
  elif path == int8_field:
    query_vector = query_text_embeddings.int8[0]
    vector_dtype = BinaryVectorDtype.INT8
  elif path == int1_field:
    query_vector = query_text_embeddings.ubinary[0]
    vector_dtype = BinaryVectorDtype.PACKED_BIT
  bson_query_vector = generate_bson_vector(query_vector, vector_dtype)
  pipeline = [
    {
      '$vectorSearch': {
        'index': index_name,
        'path': path,
        'queryVector': bson_query_vector,
        'numCandidates': <NUMBER-OF-CANDIDATES-TO-CONSIDER>, # for example, 5
        'limit': <NUMBER-OF-DOCUMENTS-TO-RETURN> # for example, 2
       }
     },
     {
       '$project': {
         '_id': 0,
         'data': 1,
         'score': { '$meta': 'vectorSearchScore' }
        }
     }
  ]
  return collection.aggregate(pipeline)

pip install --quiet --upgrade pymongo cohere

1
# Connect to your Atlas cluster
2
mongo_client = pymongo.MongoClient("<ATLAS-CONNECTION-STRING>")
3
db = mongo_client["sample_airbnb"]
4
collection = db["listingsAndReviews"]
5
6
# Filter to exclude null or empty summary fields
7
filter = { "summary": {"$nin": [None, ""]} }
8
9
# Get a subset of documents in the collection
10
documents = collection.find(filter).limit(50)
11
12
# Initialize the count of updated documents
13
updated_doc_count = 0

1
from pymongo.operations import SearchIndexModel
2
import time
3
4
# Define and create the vector search index
5
index_name = "<INDEX-NAME>"
6
search_index_model = SearchIndexModel(
7
  definition={
8
    "fields": [
9
      {
10
        "type": "vector",
11
        "path": "embedding",
12
        "similarity": "euclidean",
13
        "numDimensions": 1024
14
      }
15
    ]
16
  },
17
  name=index_name,
18
  type="vectorSearch"
19
)
20
result = collection.create_search_index(model=search_index_model)
21
print("New search index named " + result + " is building.")
22
23
# Wait for initial sync to complete
24
print("Polling to check if the index is ready. This may take up to a minute.")
25
predicate=None
26
if predicate is None:
27
  predicate = lambda index: index.get("queryable") is True
28
while True:
29
  indices = list(collection.list_search_indexes(index_name))
30
  if len(indices) and predicate(indices[0]):
31
    break
32
  time.sleep(5)
33
print(result + " is ready for querying.")