Vector Quantization

pip --quiet install 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 cohere
api_key = "<COHERE-API-KEY>"
co = cohere.Client(api_key)
generated_embeddings = co.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))

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,
          "<FIELD-NAME-FOR-FLOAT32-TYPE>":bson_f32_emb,
          "<FIELD-NAME-FOR-INT8-TYPE>":bson_int8_emb,
          "<FIELD-NAME-FOR-INT1-TYPE>":bson_int1_emb,
     }
     docs.append(doc)
  return docs
documents = create_docs_with_bson_vector_embeddings(bson_float32_embeddings, bson_int8_embeddings, bson_int1_embeddings, data)

import time
from pymongo.operations import SearchIndexModel
vector_search_index_definition = {
  "fields":[
    {
      "type": "vector",
      "path": "<FIELD-NAME-FOR-FLOAT32-TYPE>",
      "similarity": "dotProduct",
      "numDimensions": 1024,
    },
    {
      "type": "vector",
      "path": "<FIELD-NAME-FOR-INT8-TYPE>",
      "similarity": "dotProduct",
      "numDimensions": 1024,
    },
    {
      "type": "vector",
      "path": "<FIELD-NAME-FOR-INT1-TYPE>",
      "similarity": "euclidean",
      "numDimensions": 1024,
    }
  ]
}
search_index_model = SearchIndexModel(definition=vector_search_index_definition, name="<INDEX-NAME>", type="vectorSearch")
col.create_search_index(model=search_index_model)

def run_vector_search(query_text, collection, path):
  query_text_embeddings = co.embed(
    texts=[query_text],
    model="embed-english-v3.0",
    input_type="search_query",
    embedding_types=["float", "int8", "ubinary"]
  ).embeddings
  if path == "<FIELD-NAME-FOR-FLOAT32-TYPE>":
    query_vector = query_text_embeddings.float[0]
    vector_dtype = BinaryVectorDtype.FLOAT32
  elif path == "<FIELD-NAME-FOR-INT8-TYPE>":
    query_vector = query_text_embeddings.int8[0]
    vector_dtype = BinaryVectorDtype.INT8
  elif path == "<FIELD-NAME-FOR-INT1-TYPE>":
    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>,
        'limit': <NUMBER-OF-DOCUMENTS-TO-RETURN>
       }
     },
     {
       '$project': {
         '_id': 0,
         'data': 1,
         'score': { '$meta': 'vectorSearchScore' }
        }
     }
  ]
  return collection.aggregate(pipeline)

pip --quiet install 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
3
vector_search_index_definition = {
4
  "fields":[
5
    {
6
      "type": "vector",
7
      "path": "embedding",
8
      "similarity": "euclidean",
9
      "numDimensions": 1024,
10
    }
11
  ]
12
}
13
14
search_index_model = SearchIndexModel(definition=vector_search_index_definition, name="<INDEX-NAME>", type="vectorSearch")
15
16
collection.create_search_index(model=search_index_model)