efficient vector search

import faiss
import numpy as np

# Assume we have a large number of document embeddings
num_docs = 1000000
embedding_dim = 128
embeddings = np.random.random((num_docs, embedding_dim)).astype('float32')

# Create an index
index = faiss.IndexFlatL2(embedding_dim)
index.add(embeddings)

# Perform a search
k = 5  # number of nearest neighbors
query = np.random.random((1, embedding_dim)).astype('float32')
D, I = index.search(query, k)

print(f"Distances: {D}")
print(f"Indices: {I}")