Embeddings & vector databases
An embedding model turns a passage into a vector so that passages with similar meaning point in similar directions. Semantic search is then "find the stored vectors closest to the question's vector". With a few thousand vectors that is one matrix multiplication. With a hundred million you need an approximate nearest-neighbour (ANN) index, which is what a vector database is built around. This lesson covers both, with real measurements.
Finding the nearest café in a huge city
Exact search measures the distance to every café in the city. An IVF index first picks the few neighbourhoods nearest to you and only checks cafés there. HNSW is like asking for directions: start on the motorway (a sparse top layer), take the exit that heads your way, then follow smaller streets until you arrive. Both occasionally miss the true nearest café, and both are dramatically faster.
1. Similarity, concretely
| Measure | Formula | Notes |
|---|---|---|
| Cosine similarity | a·b / (|a| |b|) | the angle only; the usual choice for text |
| Dot product | a·b | identical ranking to cosine when vectors are unit length (most embedding APIs return unit vectors) |
| Euclidean (L2) | |a − b| | for unit vectors, ranks the same as cosine: |a−b|² = 2 − 2 a·b |
The examples in this module need no API key, so they use a classic dense-vector technique instead of a
neural model: LSA (TF-IDF followed by an SVD down to 256 dimensions), from
mini_rag.py. It captures word co-occurrence, not true meaning, so treat its results as a lower
bound on what text-embedding-3-small or an open model like bge or e5 would give you.
import numpy as np
from mini_rag import LsaEmbedder, doc_text
from site_docs import load_sections
docs = load_sections()
emb = LsaEmbedder([doc_text(d) for d in docs], dims=256)
print("matrix:", emb.matrix.shape, "- every row has length", round(float(np.linalg.norm(emb.matrix[0])), 3))
def search(question, k=3):
scores = emb.matrix @ emb.embed([question])[0] # cosine similarity, one matrix-vector product
for i in np.argsort(-scores)[:k]:
print(f" {scores[i]:.2f} {docs[i]['url']:44} {docs[i]['section'][:34]}")
for q in ["one partition is much bigger than the others and its task never finishes",
"retry when the API says too many requests",
"save money on model calls"]:
print(q)
search(q)
The first two land in the right neighbourhood without using the lessons' key terms. The oversized-partition
question reaches the shuffle and Spark UI stage pages (the skew lesson itself ranks sixth),
and "too many requests" finds the rate-limit sections although it never says "rate limit" or "429". The third is an honest failure. "Model calls" pulled in
the lesson on Python's callable objects (__call__), because LSA only knows which words
occur together. A neural embedding model places "save money on model calls" next to "cost optimisation".
Lesson 08 combines vectors with keyword search, and lesson 10 measures which combination works on your
questions.
2. Choosing and using an embedding model
What to compare
Retrieval quality on your own questions (lesson 10), vector size (storage and speed), maximum input length, languages, price, and whether it can run on your own hardware.
Rules that never change
Embed documents and queries with the same model. Store the model name with every vector. Changing models means re-embedding everything, because vectors from different models are not comparable.
from openai import OpenAI
client = OpenAI()
resp = client.embeddings.create(model="text-embedding-3-small",
input=[c["text"] for c in chunks[:100]], # batch: up to many inputs per call
dimensions=512) # optional: shorter vectors
vectors = [d.embedding for d in resp.data]
3. Why exact search stops scaling
Exact search costs one dot product per stored vector per query. At 1,000 vectors that is nothing. At 100 million vectors × 1,536 dimensions it is about 150 billion multiply-adds per query. ANN indexes trade a little recall (sometimes missing a true neighbour) for orders of magnitude less work.
4. Build an IVF index and measure the trade-off
Here is the whole idea in NumPy on 200,000 vectors: cluster them with k-means, keep a list of the vectors in
each cluster, and at query time scan only the nprobe closest clusters. Recall is measured
against exact search.
import time
import numpy as np
from sklearn.cluster import KMeans
rng = np.random.default_rng(0)
N, D, k = 200_000, 64, 10
centers = rng.normal(size=(50, D))
X = (centers[rng.integers(0, 50, N)] + rng.normal(size=(N, D))).astype(np.float32)
X /= np.linalg.norm(X, axis=1, keepdims=True)
Q = X[rng.choice(N, 200, replace=False)] + 0.1 * rng.normal(size=(200, D)).astype(np.float32)
Q /= np.linalg.norm(Q, axis=1, keepdims=True)
start = time.perf_counter()
truth = [set(np.argpartition(-(X @ q), k)[:k]) for q in Q] # exact top-10 for each query
print(f"exact search: 100% of vectors scanned, {(time.perf_counter() - start) / len(Q) * 1000:.2f} ms/query")
nlist = 256 # number of clusters
km = KMeans(n_clusters=nlist, n_init=1, max_iter=20, random_state=0)
km.fit(X[rng.choice(N, 20_000, replace=False)]) # train centroids on a sample
labels = km.predict(X)
order = np.argsort(labels, kind="stable") # vectors grouped by cluster
bounds = np.searchsorted(labels[order], np.arange(nlist + 1))
X_sorted, C = X[order], km.cluster_centers_.astype(np.float32)
def ivf_search(q, nprobe):
probe = np.argpartition(-(C @ q), nprobe)[:nprobe] # nearest clusters
idx = np.concatenate([np.arange(bounds[c], bounds[c + 1]) for c in probe])
top = idx[np.argpartition(-(X_sorted[idx] @ q), k)[:k]]
return set(order[top]), len(idx)
for nprobe in (1, 4, 16, 64):
start, recall, scanned = time.perf_counter(), 0.0, 0
for q, true_ids in zip(Q, truth):
found, n = ivf_search(q, nprobe)
recall += len(found & true_ids) / k
scanned += n
ms = (time.perf_counter() - start) / len(Q) * 1000
print(f"nprobe={nprobe:<3} recall@10={recall / len(Q):.2f} scanned {scanned / len(Q) / N:6.1%} {ms:.2f} ms/query")
Read the "scanned" column: scanning under 2% of the vectors already recovers about 90% of the true neighbours, and 6% recovers all of them on this data. The millisecond timings are noisy here because the Python loop dominates. Real indexes do this in C++ with SIMD, and their speed tracks the scanned fraction closely. nprobe (IVF) and ef_search (HNSW) are the dials you tune: higher means better recall and slower queries.
| Index | Strengths | Watch out for |
|---|---|---|
| Flat (exact) | perfect recall, no build step | linear cost: fine up to ~100k vectors |
| IVF (+PQ compression) | small memory with product quantisation; fast builds | must be trained; recall drops for queries near cluster borders |
| HNSW | excellent recall and speed, incremental inserts | memory-hungry graph; deletes are awkward; slower builds |
5. A real vector database, in process: Chroma
Vector databases wrap an ANN index with storage, metadata, filtering and an API. Chroma runs inside your Python process, which makes it good for prototypes. Here the vectors are passed in explicitly, so no model is downloaded, and telemetry is off:
import chromadb
from chromadb.config import Settings
from mini_rag import LsaEmbedder, doc_text
from site_docs import load_sections
docs = load_sections()
emb = LsaEmbedder([doc_text(d) for d in docs])
client = chromadb.EphemeralClient(settings=Settings(anonymized_telemetry=False))
lessons = client.create_collection("lessons", embedding_function=None,
configuration={"hnsw": {"space": "cosine"}}) # HNSW under the hood
lessons.add(ids=[d["url"] for d in docs], embeddings=emb.matrix.tolist(),
documents=[d["text"] for d in docs],
metadatas=[{"course": d["course"], "title": d["title"]} for d in docs])
print("stored:", lessons.count())
q = emb.embed(["how does an index speed up a query"])[0].tolist()
for where in (None, {"course": "sql"}):
res = lessons.query(query_embeddings=[q], n_results=3, where=where)
print("filter:", where)
for url, dist in zip(res["ids"][0], res["distances"][0]):
print(f" {1 - dist:.2f} {url}") # cosine distance = 1 - similarity
CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE chunks ( id bigserial PRIMARY KEY, course text NOT NULL, url text NOT NULL, body text NOT NULL, embedding vector(1536) NOT NULL ); CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops); -- <=> is cosine distance; the WHERE clause is an ordinary SQL filter SELECT url, 1 - (embedding <=> $1) AS similarity FROM chunks WHERE course = 'mysql' ORDER BY embedding <=> $1 LIMIT 5;
from qdrant_client import QdrantClient, models
client = QdrantClient(url="http://localhost:6333")
client.create_collection("lessons", vectors_config=models.VectorParams(size=1536, distance=models.Distance.COSINE))
client.upsert("lessons", points=[models.PointStruct(id=i, vector=v, payload=meta)
for i, (v, meta) in enumerate(zip(vectors, metadatas))])
hits = client.query_points("lessons", query=query_vector, limit=5,
query_filter=models.Filter(must=[models.FieldCondition(
key="course", match=models.MatchValue(value="mysql"))])).points
import faiss import numpy as np index = faiss.IndexHNSWFlat(1536, 32, faiss.METRIC_INNER_PRODUCT) # 32 links per node index.add(np.asarray(vectors, dtype="float32")) # unit vectors: inner product = cosine index.hnsw.efSearch = 64 # the recall/speed dial scores, ids = index.search(np.asarray([query_vector], dtype="float32"), 5) # metadata, filtering and persistence are your job
The pgvector, Qdrant and FAISS snippets follow those projects' documented APIs but were not run here (they need a server or a package this site does not install). The Chroma example above was run.
6. Which one should you use?
| Situation | Good default |
|---|---|
| Under ~100k chunks, one process | a NumPy matrix, exactly as in section 1. No infrastructure. |
| You already run Postgres | pgvector: one database, transactions, joins and SQL filters |
| Prototype or local tool | Chroma (in-process or client/server) |
| Tens of millions of vectors, heavy filtering, many tenants | Qdrant, Weaviate or Milvus, or a managed service |
| Already on a platform | its native option: Databricks Vector Search, OpenAI vector stores, a cloud provider's search service |
Recap
- Embeddings turn text into vectors; cosine similarity ranks them. Use one model per index, and store its name.
- Exact search is fine up to ~100k vectors. Beyond that, ANN (IVF, HNSW) trades a little recall for big speed-ups.
- Tune nprobe / ef_search against measured recall, not by guessing.
- Pick the simplest store that meets your scale: NumPy → pgvector/Chroma → a dedicated vector DB.