Runs in your browser — no install. (Works once the repo is public.)

In [ ]:

# --- Colab setup (skipped when running locally) ---
import os, sys
if 'google.colab' in sys.modules and not os.path.exists('data/sample_corpus.json'):
    !git clone -q https://github.com/mdhabibi/llm-search-handbook.git
    %cd llm-search-handbook
    !pip -q install -r requirements.txt

# --- Colab setup (skipped when running locally) --- import os, sys if 'google.colab' in sys.modules and not os.path.exists('data/sample_corpus.json'): !git clone -q https://github.com/mdhabibi/llm-search-handbook.git %cd llm-search-handbook !pip -q install -r requirements.txt

Chapter 6 — Vector Databases & ANN¶

Brute force is fine for 16 passages, hopeless for 16 million. Here we make nearest-neighbor search fast at scale with FAISS (exact + HNSW) and feel the recall/speed trade-off. The FAISS parts run fully offline on synthetic vectors.

Setup¶

pip install faiss-cpu chromadb sentence-transformers

In [ ]:

# Bootstrap: locate repo root and import shared helpers
import sys, os
d = os.getcwd()
while d != os.path.dirname(d) and not os.path.exists(os.path.join(d, 'data', 'sample_corpus.json')):
    d = os.path.dirname(d)
ROOT = d; sys.path.insert(0, os.path.join(ROOT, 'src'))
import numpy as np
from corpus import load_corpus, tokenize

# Bootstrap: locate repo root and import shared helpers import sys, os d = os.getcwd() while d != os.path.dirname(d) and not os.path.exists(os.path.join(d, 'data', 'sample_corpus.json')): d = os.path.dirname(d) ROOT = d; sys.path.insert(0, os.path.join(ROOT, 'src')) import numpy as np from corpus import load_corpus, tokenize

In [ ]:

import faiss, time

import faiss, time

1. Exact search with FAISS (matches Chapter 5's brute force)¶

IndexFlatIP does exact inner-product search. On L2-normalized vectors, inner product = cosine.

In [ ]:

rng = np.random.default_rng(0)
dim = 128
data = rng.standard_normal((2000, dim)).astype('float32')
faiss.normalize_L2(data)                     # so inner product == cosine
query = data[:5].copy()                      # use 5 real vectors as queries

flat = faiss.IndexFlatIP(dim); flat.add(data)
D, I = flat.search(query, 5)
print('FAISS exact top-5 for query 0:', I[0])

# Cross-check against pure-numpy brute force
brute = np.argsort(-(data @ query[0]))[:5]
print('numpy brute force        :', brute)
print('match?', list(I[0]) == list(brute))

rng = np.random.default_rng(0) dim = 128 data = rng.standard_normal((2000, dim)).astype('float32') faiss.normalize_L2(data) # so inner product == cosine query = data[:5].copy() # use 5 real vectors as queries flat = faiss.IndexFlatIP(dim); flat.add(data) D, I = flat.search(query, 5) print('FAISS exact top-5 for query 0:', I[0]) # Cross-check against pure-numpy brute force brute = np.argsort(-(data @ query[0]))[:5] print('numpy brute force :', brute) print('match?', list(I[0]) == list(brute))

2. Approximate search with HNSW¶

HNSW builds a navigable graph so a query only visits a small neighborhood instead of all N vectors. We build one over a larger set.

Note on the test data: pure random vectors are a worst case for ANN (no structure to exploit). Real embeddings cluster by meaning, which is exactly what ANN thrives on — so we add mild cluster structure to make this synthetic benchmark representative.

In [ ]:

# 100k vectors with mild cluster structure (like real embeddings clustering by topic)
centers = rng.standard_normal((50, dim)).astype('float32')
assign = rng.integers(0, 50, size=100000)
big = (centers[assign] + 0.6 * rng.standard_normal((100000, dim))).astype('float32')
faiss.normalize_L2(big)
queries = big[:1000].copy()

hnsw = faiss.IndexHNSWFlat(dim, 32)          # 32 = graph connectivity (M)
hnsw.hnsw.efConstruction = 200
hnsw.hnsw.efSearch = 64                       # search-time breadth (recall knob)
hnsw.add(big)
print('HNSW index built over', hnsw.ntotal, 'vectors.')

# 100k vectors with mild cluster structure (like real embeddings clustering by topic) centers = rng.standard_normal((50, dim)).astype('float32') assign = rng.integers(0, 50, size=100000) big = (centers[assign] + 0.6 * rng.standard_normal((100000, dim))).astype('float32') faiss.normalize_L2(big) queries = big[:1000].copy() hnsw = faiss.IndexHNSWFlat(dim, 32) # 32 = graph connectivity (M) hnsw.hnsw.efConstruction = 200 hnsw.hnsw.efSearch = 64 # search-time breadth (recall knob) hnsw.add(big) print('HNSW index built over', hnsw.ntotal, 'vectors.')

3. Benchmark: recall@10 and speed vs. exact¶

The key experiment. We compare HNSW to exact search on the same 1000 queries: how much faster, and how much recall do we lose?

In [ ]:

exact = faiss.IndexFlatIP(dim); exact.add(big)

t0 = time.time(); _, I_exact = exact.search(queries, 10); t_exact = time.time()-t0
t0 = time.time(); _, I_hnsw  = hnsw.search(queries, 10);  t_hnsw  = time.time()-t0

# recall@10 = fraction of each query's true top-10 that HNSW also returned
recall = np.mean([len(set(a) & set(b))/10 for a, b in zip(I_exact, I_hnsw)])
print(f'exact search: {t_exact*1000:7.1f} ms for 1000 queries')
print(f'HNSW  search: {t_hnsw*1000:7.1f} ms for 1000 queries  ({t_exact/max(t_hnsw,1e-9):.1f}x faster)')
print(f'recall@10   : {recall:.3f}   (fraction of exact neighbors recovered)')
print('\nRaise efSearch -> higher recall, slower. Lower -> faster, less recall. That is the trade-off.')

exact = faiss.IndexFlatIP(dim); exact.add(big) t0 = time.time(); _, I_exact = exact.search(queries, 10); t_exact = time.time()-t0 t0 = time.time(); _, I_hnsw = hnsw.search(queries, 10); t_hnsw = time.time()-t0 # recall@10 = fraction of each query's true top-10 that HNSW also returned recall = np.mean([len(set(a) & set(b))/10 for a, b in zip(I_exact, I_hnsw)]) print(f'exact search: {t_exact*1000:7.1f} ms for 1000 queries') print(f'HNSW search: {t_hnsw*1000:7.1f} ms for 1000 queries ({t_exact/max(t_hnsw,1e-9):.1f}x faster)') print(f'recall@10 : {recall:.3f} (fraction of exact neighbors recovered)') print('\nRaise efSearch -> higher recall, slower. Lower -> faster, less recall. That is the trade-off.')

4. A real vector database: Chroma with persistence + metadata filtering¶

An index alone isn't a product. Chroma adds persistence and metadata filters. Here we store our course passages with a language field and run a filtered semantic query. (Uses the embedding model, so runs on your machine after the first download.)

In [ ]:

import chromadb
from sentence_transformers import SentenceTransformer

docs = load_corpus()
model = SentenceTransformer('all-MiniLM-L6-v2')
embs = model.encode([d['text'] for d in docs]).tolist()

client = chromadb.Client()                    # in-memory; use PersistentClient(path=...) to save
col = client.get_or_create_collection('course_corpus')
col.add(
    ids=[str(d['id']) for d in docs],
    embeddings=embs,
    documents=[d['text'] for d in docs],
    metadatas=[{'title': d['title'], 'language': 'en'} for d in docs],
)

q = model.encode(['a large ocean animal']).tolist()
res = col.query(query_embeddings=q, n_results=3, where={'language': 'en'})
print('Filtered semantic search (language == en):')
for doc, meta in zip(res['documents'][0], res['metadatas'][0]):
    print('  -', meta['title'])

import chromadb from sentence_transformers import SentenceTransformer docs = load_corpus() model = SentenceTransformer('all-MiniLM-L6-v2') embs = model.encode([d['text'] for d in docs]).tolist() client = chromadb.Client() # in-memory; use PersistentClient(path=...) to save col = client.get_or_create_collection('course_corpus') col.add( ids=[str(d['id']) for d in docs], embeddings=embs, documents=[d['text'] for d in docs], metadatas=[{'title': d['title'], 'language': 'en'} for d in docs], ) q = model.encode(['a large ocean animal']).tolist() res = col.query(query_embeddings=q, n_results=3, where={'language': 'en'}) print('Filtered semantic search (language == en):') for doc, meta in zip(res['documents'][0], res['metadatas'][0]): print(' -', meta['title'])

Takeaway & exercises¶

ANN trades a sliver of recall for enormous speed; a vector DB wraps that with persistence, metadata, and CRUD. Choose based on corpus size — under ~100k vectors, exact is often fine.

Exercises

1. Sweep efSearch in {16, 32, 64, 128}. Plot recall@10 vs. query time.
2. Grow the corpus to 200k vectors. How does exact search time scale vs. HNSW?
3. Switch Chroma to PersistentClient(path='.chroma'), restart the kernel, and reload without re-embedding.
4. Add a year metadata field and filter where={'year': {'$gte': 2020}}.