Vector Search with Kai

1!pip install datasets --quiet

1import os2import pprint3import time4import concurrent.futures5import datasets6from pymongo import MongoClient7from datasets import load_dataset8from bson import json_util

1current_database = %sql SELECT DATABASE() as CurrentDatabase2DB = current_database[0][0]3COLLECTION = 'wiki_embeddings'

1# Using the environment variable that holds the kai endpoint2client = MongoClient(connection_url_kai)3collection = client[DB][COLLECTION]

1client[DB].create_collection(COLLECTION,2  columns=[{ 'id': "emb", 'type': "VECTOR(768) NOT NULL" }],3);

1# Using the "wikipedia-22-12-simple-embeddings" dataset from Hugging Face2dataset = load_dataset("Cohere/wikipedia-22-12-simple-embeddings", split="train")

1DB_SIZE = 50000 #Currently loading 50k documents to the collection, can go to a max of 485,859 for this dataset2insert_data = []3insert_count = 04# Iterate through the dataset and prepare the documents for insertion5# The script below ingests 1000 records into the database at a time6for item in dataset:7    if insert_count >= DB_SIZE:8        break9    # Convert the dataset item to MongoDB document format10    doc_item = json_util.loads(json_util.dumps(item))11    insert_data.append(doc_item)12
13    # Insert in batches of 1000 documents14    if len(insert_data) == 1000:15        collection.insert_many(insert_data)16        insert_count += 100017        print(f"{insert_count} of {DB_SIZE} records ingested")18        insert_data = []19
20
21# Insert any remaining documents22if len(insert_data) > 0:23    collection.insert_many(insert_data)24    print("Data Ingested")

1sample_doc = collection.find_one()2pprint.pprint(sample_doc, compact=True)

1client[DB].command({2    'createIndexes': COLLECTION,3    'indexes': [{4        'key': {'emb': 'vector'},5        'name': 'vector_index',6        'kaiSearchOptions': {"index_type":"AUTO", "metric_type": "EUCLIDEAN_DISTANCE", "dimensions": 768}7    }],8})

1# input vector2query_vector = sample_doc['emb']

1def execute_kai_search(query_vector):2    pipeline = [3        {4            '$vectorSearch': {5                "index": "vector_index",6                "path": "emb",7                "queryVector": query_vector,8                "numCandidates": 20,9                "limit": 3,10            }11        },12        {13            '$project': {14               '_id':1,15               'text': 1,16            }17        }18    ]19    results = collection.aggregate(pipeline)20    return list(results)

1execute_kai_search(query_vector)

1num_concurrent_queries = 2502start_time = time.time()3
4with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_queries) as executor:5    futures = [executor.submit(execute_kai_search, query_vector) for _ in range(num_concurrent_queries)]6    concurrent.futures.wait(futures)7
8end_time = time.time()9print(f"Executed {num_concurrent_queries} concurrent queries.")10print(f"Total execution time: {end_time - start_time} seconds")11
12for f in futures:13    if f.exception() is not None:14        print(f.exception())15failed_count = sum(1 for f in futures if f.exception() is not None)16print(f"Failed queries: {failed_count}")