Semantic Search

from sentence_transformers import SentenceTransformer
import numpy as np

# Load pre-trained model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Encode documents
documents = [
    "Machine learning is a subset of artificial intelligence.",
    "Deep learning uses neural networks with many layers.",
    "Natural language processing deals with text and speech.",
    "Computer vision analyzes images and videos.",
]

# Generate embeddings
embeddings = model.encode(documents)
print(f"Embedding shape: {embeddings.shape}")  # (4, 384)

# Compute similarity
query = "What is AI?"
query_embedding = model.encode([query])

similarities = np.dot(embeddings, query_embedding.T).flatten()
ranked_indices = np.argsort(similarities)[::-1]

for idx in ranked_indices:
    print(f"Score: {similarities[idx]:.4f} | {documents[idx]}")

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# Batch computation
def batch_cosine_similarity(query_emb, doc_embs):
    # Normalize
    query_norm = query_emb / np.linalg.norm(query_emb, axis=1, keepdims=True)
    doc_norm = doc_embs / np.linalg.norm(doc_embs, axis=1, keepdims=True)
    return np.dot(doc_norm, query_norm.T)

import faiss
import numpy as np

# Generate random embeddings for demonstration
dimension = 384
num_documents = 100000
np.random.seed(42)

documents = np.random.random((num_documents, dimension)).astype('float32')

# Build FAISS index
# Option 1: Flat index (exact search)
index_flat = faiss.IndexFlatL2(dimension)
index_flat.add(documents)

# Option 2: IVF index (approximate search)
nlist = 100  # Number of clusters
quantizer = faiss.IndexFlatL2(dimension)
index_ivf = faiss.IndexIVFFlat(quantizer, dimension, nlist)
index_ivf.train(documents)
index_ivf.add(documents)

# Option 3: HNSW index (graph-based)
index_hnsw = faiss.IndexHNSWFlat(dimension, 32)
index_hnsw.hnsw.efConstruction = 64
index_hnsw.add(documents)

# Search
query = np.random.random((1, dimension)).astype('float32')
k = 5  # Number of nearest neighbors

# Flat search (exact)
distances, indices = index_flat.search(query, k)
print(f"Flat results: {indices[0]}, distances: {distances[0]}")

# IVF search (approximate)
index_ivf.nprobe = 10  # Number of clusters to search
distances, indices = index_ivf.search(query, k)
print(f"IVF results: {indices[0]}, distances: {distances[0]}")

# HNSW search
distances, indices = index_hnsw.search(query, k)
print(f"HNSW results: {indices[0]}, distances: {distances[0]}")

def select_index(num_vectors, dimension, speed_priority=True):
    if num_vectors < 10000:
        return "IndexFlatL2"  # Exact search
    elif num_vectors < 1000000:
        if speed_priority:
            return "IndexHNSWFlat"  # Fast, good accuracy
        else:
            return "IndexIVFFlat"  # Balanced
    else:
        if speed_priority:
            return "IndexIVFPQ"  # Memory efficient
        else:
            return "IndexHNSWFlat"  # Good accuracy

# Example usage
index_type = select_index(500000, 768, speed_priority=True)
print(f"Recommended index: {index_type}")

# Example with ChromaDB
import chromadb

client = chromadb.Client()
collection = client.create_collection("documents")

# Add documents
documents = [
    "Machine learning algorithms learn patterns from data.",
    "Deep neural networks have multiple hidden layers.",
    "Natural language processing understands human language.",
]
ids = ["doc1", "doc2", "doc3"]

collection.add(
    documents=documents,
    ids=ids,
    metadatas=[{"source": "textbook"} for _ in documents]
)

# Query
results = collection.query(
    query_texts=["What is deep learning?"],
    n_results=2
)
print(results["documents"])