RAG: read before answering

TF-IDF and cosine similarity

This is how the scene’s word search is computed. First, very common words are removed (“the,” “of,” “and”…) and each word is cut down to its stem in a simple way (“batteries” and “battery” count the same). With that, the lesson’s manual has 228 distinct terms: each chunk is a vector of 228 numbers.

Given $N$ chunks, $c_{t,d}$ is how many times term $t$ appears in chunk $d$, and $\mathrm{df}(t)$ is how many chunks it appears in. With the variants used in the lesson:

With $N=25$: “cycles” appears in 1 chunk ($\mathrm{idf}\approx 3.56$) and “warranty” in 6 ($\mathrm{idf}\approx 2.31$); that is why “cycles” weighs more. Each vector is normalized, $\hat{\mathbf d}=\mathbf d/\lVert\mathbf d\rVert$, and the question gets the same treatment with the index’s $\mathrm{idf}$ values. Cosine similarity becomes a dot product:

Weighting by rarity goes back to Karen Spärck Jones (1972): a term that appears in few documents tells them apart better.

BM25, the other word-based score

BM25 (Robertson and Zaragoza, 2009) scores each chunk $D$ with a saturating term frequency and a length correction:

with $1.2<k_1<2$ and $0.5<b<0.8$ as values that, according to its authors, work reasonably well in many cases. It is a hard baseline to beat with words alone, and many systems combine it with embeddings (hybrid search).

Embeddings and approximate search

With embeddings (lesson 14) the vector has a few hundred dense numbers and the words do not need to match: cosine similarity compares meanings. Comparing the question with all $N$ vectors costs $O(N\cdot D)$; with millions of chunks, approximate indexes such as HNSW (Malkov and Yashunin) trade a little exactness for a lot of speed. In the paper that named RAG (Lewis et al., 2020), Wikipedia was cut into 100-word pieces, 21 million in all, indexed with FAISS and HNSW; during training, $k=5$ or $k=10$ were retrieved per question.

Reranking

A dense retriever (bi-encoder), such as DPR (Karpukhin et al., 2020), the one RAG uses, encodes the question and the chunk separately, so chunks are turned into vectors only once. A reranker (cross-encoder) reads the question and the chunk together in one model pass and returns a score: it is finer, but it costs one pass per candidate, so it is only applied to the top of the list. Nogueira and Cho (2019) used BERT to rerank the top 1,000 passages returned by BM25.

In the lesson, the embedding and reranker figures were measured offline (see build-dense.py) with two open models: intfloat/multilingual-e5-small (Wang et al., 2024) and cross-encoder/mmarco-mMiniLMv2-L12-H384-v1, trained on mMARCO (Bonifacio et al., 2021). The reranker is applied to the top 5 candidates.

How retrieval is measured

For each test question you record the right chunk and its position $r_i$ in the ranked list ($r_i=0$ if it does not appear). With $Q$ questions:

On the scene’s 10 questions, with $k=3$: word search, 7 of 10; embeddings, 9 of 10; embeddings plus reranker, 8 of 10 (in both languages). Ten questions are far too few for general conclusions: the point is the method.

This measures only the first half. The second, whether the answer relies on what was retrieved, is evaluated with test cases and with human review or a model as judge (lesson 34). Ragas (Es et al., 2023) splits it into three measures: faithfulness (what the answer says follows from the context), answer relevance, and context relevance.

Updating without retraining

The RAG paper tested swapping only the index: for 82 world leaders who changed between 2016 and 2018, the model was right 70% of the time with the 2016 Wikipedia index for 2016 leaders, and 68% with the 2018 index for 2018 leaders; with the mismatched index, only 12% and 4%. To update what it knows, you change documents, not weights.

In code

A complete TF-IDF search, with no libraries. It is the same recipe as the lesson, without stemming:

import math, re
from collections import Counter

STOP = {"the", "has", "an", "or", "you", "can", "your", "of", "on", "is", "at", "how", "long"}

def terms(text):
    return [w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOP]

chunks = [
    "the battery has an 18 month warranty or 500 charge cycles",
    "you can cancel your purchase within 30 days of delivery",
    "shipping is free on purchases over 1500",
    "charge the battery at room temperature",
]

N = len(chunks)
bags = [Counter(terms(c)) for c in chunks]
df = Counter(t for b in bags for t in b)
idf = {t: math.log((N + 1) / (df[t] + 1)) + 1 for t in df}

def vector(bag):
    v = {t: (1 + math.log(c)) * idf[t] for t, c in bag.items() if t in idf}
    norm = math.sqrt(sum(x * x for x in v.values())) or 1.0
    return {t: x / norm for t, x in v.items()}

index = [vector(b) for b in bags]

def search(question, k=2):
    q = vector(Counter(terms(question)))
    scores = [sum(w * d.get(t, 0.0) for t, w in q.items()) for d in index]
    return sorted(enumerate(scores), key=lambda p: -p[1])[:k]

for i, s in search("how long is the battery warranty"):
    print(f"{s:.2f}  {chunks[i]}")
# 0.51  the battery has an 18 month warranty or 500 charge cycles
# 0.27  charge the battery at room temperature

What comes next

Here the model only reads. In lesson 32 it can also use tools: search, calculate, write. Prompt injection comes back in lesson 35, and full system evaluation in lesson 34.