Como ingestão de vetores quantizados

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"]
).embeddings
float32_embeddings = generated_embeddings.float
int8_embeddings = generated_embeddings.int8

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 and int8 embeddings
bson_float32_embeddings = []
bson_int8_embeddings = []
for i, (f32_emb, int8_emb) in enumerate(zip(float32_embeddings, int8_embeddings)):
   bson_float32_embeddings.append(generate_bson_vector(f32_emb, BinaryVectorDtype.FLOAT32))
   bson_int8_embeddings.append(generate_bson_vector(int8_emb, BinaryVectorDtype.INT8))

def create_docs_with_bson_vector_embeddings(bson_float32_embeddings, bson_int8_embeddings, data):
  docs = []
  for i, (bson_f32_emb, bson_int8_emb, text) in enumerate(zip(bson_float32_embeddings, bson_int8_embeddings, data)):
     doc = {
          "_id":i,
          "data": text,
          "<FIELD-NAME-FOR-INT8-TYPE>":bson_int8_emb,
          "<FIELD-NAME-FOR-FLOAT32-TYPE>":bson_f32_emb,
     }
     docs.append(doc)
  return docs
documents = create_docs_with_bson_vector_embeddings(bson_float32_embeddings, bson_int8_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,
    }
  ]
}
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"]
  ).embeddings
  if path == "<FIELD-NAME-FOR-FLOAT32-TYPE>":
    query_vector = query_text_embeddings.float[0]
    vector_dtype = BinaryVectorDtype.FLOAT32
  else:
    query_vector = query_text_embeddings.int8[0]
    vector_dtype = BinaryVectorDtype.INT8
  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
import os
2
import pymongo
3
import cohere
4
from bson.binary import Binary, BinaryVectorDtype
5
6
# Specify your OpenAI API key and embedding model
7
os.environ["COHERE_API_KEY"] = "<COHERE-API-KEY>"
8
cohere_client = cohere.Client(os.environ["COHERE_API_KEY"])
9
10
# Function to generate embeddings using Cohere
11
def get_embedding(text):
12
    response = cohere_client.embed(
13
      texts=[text],
14
      model='embed-english-v3.0',
15
      input_type='search_document'
16
    )
17
    embedding = response.embeddings[0]
18
    return embedding
19
20
# Function to convert embeddings to BSON-compatible format
21
def generate_bson_vector(vector, vector_dtype):
22
    return Binary.from_vector(vector, vector_dtype)

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
for doc in documents:
2
    # Generate embeddings based on the summary
3
    summary = doc["summary"]
4
    embedding = get_embedding(summary)  # Get float32 embedding
5
6
    # Convert the float32 embedding to BSON format
7
    bson_float32 = generate_bson_vector(embedding, BinaryVectorDtype.FLOAT32)
8
9
    # Update the document with the BSON embedding
10
    collection.update_one(
11
        {"_id": doc["_id"]},
12
        {"$set": {"embedding": bson_float32}}
13
    )
14
    updated_doc_count += 1
15
16
print(f"Updated {updated_doc_count} documents with BSON embeddings.")

1
from pymongo.operations import SearchIndexModel
2
3
vector_search_index_definition = {
4
  "fields":[
5
    {
6
      "type": "vector",
7
      "path": "embedding",
8
      "similarity": "dotProduct",
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)

1
def run_vector_search(query_text, collection, path):
2
  query_embedding = get_embedding("query_text")
3
  bson_query_vector = generate_bson_vector(query_embedding, BinaryVectorDtype.FLOAT32)
4
5
  pipeline = [
6
    {
7
      '$vectorSearch': {
8
        'index': '<INDEX-NAME>',
9
        'path': path,
10
        'queryVector': bson_query_vector,
11
        'numCandidates': <NUMBER-OF-CANDIDATES-TO-CONSIDER>, # for example, 20
12
        'limit': <NUMBER-OF-DOCUMENTS-TO-RETURN> # for example, 5
13
       }
14
     },
15
     {
16
       '$project': {
17
         '_id': 0,
18
         'name': 1,
19
         'summary': 1,
20
         'score': { '$meta': 'vectorSearchScore' }
21
        }
22
     }
23
  ]
24
25
  return collection.aggregate(pipeline)