AI Engineering: 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 part covers both, with real measurements.

1. Similarity, concretely

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.

semantic search over this site's 1,229 part sections

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)

Output

matrix: (1229, 256) - every row has length 1.0
one partition is much bigger than the others and its task never finishes
  0.47  course/19-shuffle.html#anatomy                1. What actually happens during a 
  0.41  course/02-architecture.html#vocab             3. Seven words you must not mix up
  0.41  course/22-spark-ui.html#stage                 3. The stage detail page — the fiv
retry when the API says too many requests
  0.39  openai-sdk/11-async-concurrency.html          Introduction
  0.29  openai-sdk/12-errors-retries.html             Introduction
  0.29  openai-sdk/11-async-concurrency.html#limits   3. Rate limits: RPM and TPM
save money on model calls
  0.43  python/09-dunder-methods.html#call            4. Callable objects
  0.37  python/09-dunder-methods.html                 Introduction
  0.35  openai-sdk/13-cost-latency.html#latency       5. Latency checklist

The first two land in the right neighbourhood without using the parts’ key terms. The oversized-partition question reaches the shuffle and Spark UI stage pages (the skew part 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 part 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”. Part 08 combines vectors with keyword search, and part 10 measures which combination works on your questions.

2. Choosing and using an embedding model

the real call (OpenAI SDK guide, part 10)

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.

IVF: recall vs work, measured

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")

Output

exact search: 100% of vectors scanned, 3.41 ms/query
nprobe=1   recall@10=0.40  scanned   0.4%  0.09 ms/query
nprobe=4   recall@10=0.90  scanned   1.7%  0.30 ms/query
nprobe=16  recall@10=1.00  scanned   6.4%  2.13 ms/query
nprobe=64  recall@10=1.00  scanned  25.3%  6.52 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.

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:

Chroma with our own vectors and a metadata filter

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

Output

stored: 1229
filter: None
   0.48  mysql/12-covering-indexes.html#covering
   0.40  mysql/12-covering-indexes.html#recap
   0.40  sql/26-indexes.html#covering
filter: {'course': 'sql'}
   0.40  sql/26-indexes.html#covering
   0.38  sql/27-explain.html#sqlite
   0.35  sql/26-indexes.html#measured

pgvector

Postgres + pgvector: vectors next to your relational data

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;

Qdrant

Qdrant: a dedicated vector database (pip install qdrant-client)

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

FAISS

FAISS: an index library, not a database (pip install faiss-cpu)

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?

Recap