Module 15: Retrieval-Augmented Generation

Introduction

Your model (m06–m08) knows exactly what was in its training data, frozen at a cutoff date, blended into billions of weights. Ask it about a document it never saw — your company’s wiki, a paper from last week, a fact it half-remembers — and it will confidently make something up. Retrieval-Augmented Generation (RAG) fixes this without touching the weights: at query time, search an external corpus for relevant passages and paste them into the prompt. The model then answers from fresh, specific, checkable text instead of parametric memory alone.

The engine underneath is nearest-neighbor search in an embedding space. Encode every document as a vector; encode the question the same way; retrieve the documents whose vectors point most nearly the same direction (highest cosine similarity); stuff them into the prompt; generate. This module builds that whole pipeline from scratch.

To keep it fully runnable with no trained encoder, we encode with classic TF-IDF vectors — term frequency times inverse document frequency, computable by hand. Modern RAG swaps TF-IDF for dense embeddings from a trained encoder (the m04 idea, learned end-to-end), but the pipeline is identical — only the encode step changes.

Why it matters for LLMs:

  • Knowledge without retraining. New facts arrive by adding documents to the index, not by a training run. The model stays fixed; the corpus is live.
  • Grounding and citations. The answer is conditioned on retrieved text you can show the user — the standard defense against hallucination.
  • It is everywhere. Chat-with-your-docs, coding assistants over a repo, search copilots — almost every applied LLM system is a RAG system.

What You’ll Learn

After this module, you can:

  • Explain the RAG pipeline — encode → search → retrieve → augment → generate — and why it beats parametric memory for fresh or private knowledge.
  • Build TF-IDF document vectors from scratch (tf, smoothed idf).
  • Rank documents by cosine similarity and retrieve the top-k.
  • Chunk long documents, and de-duplicate results with Maximal Marginal Relevance (MMR).
  • Assemble a grounded prompt and wrap it all in a Retriever.

Prerequisites

This module requires familiarity with:

  • Module 04: Embeddings — vectors as meaning; dense retrieval replaces TF-IDF with learned embeddings of exactly this kind.
  • Module 08: Generation — RAG conditions generation on retrieved context; the decoding is unchanged.

Intuition: The RAG Pipeline

RAG is five steps, and only the middle three are new — the model at the end is the same one you already built. Step through what each stage produces:

NoteKey Insight

Only Encode → Search → Retrieve → Augment is retrieval; Generate is the model you already have. RAG is not a new kind of model — it is a way of choosing what goes in the prompt. Everything hard is in ranking documents by relevance.

The Math: TF-IDF and Cosine Similarity

To search by meaning we need documents as vectors. The classic recipe weights each term by how often it appears in a document (term frequency) against how rare it is across the corpus (inverse document frequency), so shared rare words count and boilerplate like “the” does not:

\[ \text{tf}(t, d) = \text{count of } t \text{ in } d, \qquad \text{idf}(t) = \ln\!\frac{1 + N}{1 + \text{df}(t)} + 1, \qquad \text{tfidf}(t, d) = \text{tf}(t, d)\cdot\text{idf}(t), \]

where \(N\) is the number of documents and \(\text{df}(t)\) how many contain \(t\). The \(+1\) smoothing keeps every idf positive; a term in every document gets idf \(= 1\).

Relevance is the angle between vectors, not their length — a long document is not more relevant just for being long — so we compare with cosine similarity:

\[ \cos(\mathbf{q}, \mathbf{d}) = \frac{\mathbf{q}\cdot\mathbf{d}}{\lVert\mathbf{q}\rVert\,\lVert\mathbf{d}\rVert}. \]

Normalize both vectors to unit length and cosine is just a dot product. retrieval.py implements tfidf_matrix, cosine_similarity, and retrieve.

Code: Build a Retriever

The Retriever class fits a corpus once, then answers queries by cosine top-k:

from retrieval import Retriever

corpus = [
    "The transformer architecture uses self-attention to relate every token to every other token.",
    "Attention computes a weighted sum of value vectors using query-key similarity scores.",
    "Byte-pair encoding builds a subword vocabulary by merging the most frequent adjacent pairs.",
    "Photosynthesis lets plants convert sunlight, water, and carbon dioxide into glucose and oxygen.",
    "Mount Everest is the tallest mountain above sea level, on the border of Nepal and Tibet.",
]

retriever = Retriever().fit(corpus)
for idx, score, doc in retriever.search("how does attention work?", k=2):
    print(f"[{idx}] {score:.3f}  {doc[:60]}...")
[1] 0.233  Attention computes a weighted sum of value vectors using que...
[0] 0.185  The transformer architecture uses self-attention to relate e...

The attention documents win because they share the rare, informative words of the query. Now assemble the retrieved passages into a grounded prompt — the “augment” step:

from retrieval import build_prompt

query = "how does attention work?"
hits = retriever.search(query, k=2)
prompt = build_prompt(query, [doc for _, _, doc in hits])
print(prompt)
Answer the question using only the context below. If the answer is not in the context, say you don't know.

Context:
[1] Attention computes a weighted sum of value vectors using query-key similarity scores.
[2] The transformer architecture uses self-attention to relate every token to every other token.

Question: how does attention work?
Answer:

That string is exactly what you would hand to generate() from m08. The model now answers from the retrieved context — and because you have the passages, you can show them as citations. Swap the TF-IDF encode for a trained dense encoder and nothing else in this pipeline changes.

Interactive: Watch Retrieval Rank the Corpus

Pick a question and see cosine similarity score every document. The top-k (here k = 3) are what gets retrieved and pasted into the prompt; everything else is ignored. Notice how the scores concentrate on the documents that share the query’s informative words.

TipTry This
  1. Switch between the attention query and the photosynthesis query. Watch the mass of similarity jump to a completely different pair of documents — the corpus is the same; only the query moved.
  2. Note the ignored documents still get a small nonzero score from incidental shared words (“the”, “of”). idf is what keeps those from dominating.

Chunking and Diversity

Two practical problems break naïve retrieval, and retrieval.py handles both.

Chunking. You do not index whole documents — a 50-page PDF has one vector that means nothing specific. You split it into small overlapping windows so each retrievable passage is focused, with a little overlap so a fact spanning a boundary survives in at least one chunk:

from retrieval import chunk_text

passage = "Retrieval augmented generation searches a corpus then conditions the model on the results"
print(chunk_text(passage, chunk_size=5, overlap=2))
['Retrieval augmented generation searches a', 'searches a corpus then conditions', 'then conditions the model on', 'model on the results']

Diversity. Plain top-k happily returns three near-identical passages, wasting the context window. Maximal Marginal Relevance (MMR) picks each next passage for relevance minus similarity to what is already chosen:

\[ \text{next} = \arg\max_{c}\;\Big[\lambda\cdot\text{rel}(c) - (1-\lambda)\max_{s\in S}\text{sim}(c, s)\Big]. \]

from retrieval import tfidf_matrix, encode_query, retrieve, mmr

docs = [
    "transformer attention model",
    "transformer attention model network",   # near-duplicate of doc 0
    "retrieval augmented generation search",
]
matrix, vocab, idf = tfidf_matrix(docs)
qv = encode_query("transformer attention", vocab, idf)

print("plain top-2:", [i for i, _ in retrieve(qv, matrix, k=2)])   # two near-duplicates
print("MMR top-2:  ", mmr(qv, matrix, k=2, lambda_=0.3))            # swaps in the different doc
plain top-2: [0, 1]
MMR top-2:   [0, 2]

With \(\lambda = 1\), MMR is just top-k; lower \(\lambda\) trades a little relevance for coverage.

Common Pitfalls

  1. Chunk too big or too small. Whole-document chunks retrieve nothing specific; one-sentence chunks lose context. A few hundred words with light overlap is the usual sweet spot.
  2. Lexical retrieval misses paraphrases. TF-IDF matches words, so “car” won’t retrieve a passage about “automobiles.” This is exactly why production RAG uses dense embeddings — they match meaning, not surface form.
  3. Retrieving duplicates. Top-k over a corpus with repeats fills the context with the same fact. De-duplicate (MMR) or the model sees no new information.
  4. Stuffing too much context. More passages is not better — irrelevant context distracts the model and costs tokens. Retrieve few, retrieve well.
  5. Trusting retrieval blindly. If the answer isn’t in the corpus, a grounded prompt should let the model say so — the build_prompt instruction asks for exactly that. RAG reduces hallucination; it does not abolish it.

Exercises

Exercise 1: idf by hand

from retrieval import build_vocab, compute_idf

# For docs ["the cat sat", "the dog ran", "the cat ran"], compute idf by hand
# (N=3) for "the" (df=3), "cat" (df=2), "sat" (df=1). Confirm compute_idf agrees
# and that the rarest term has the highest idf.

# Your implementation here:

Exercise 2: retrieve then augment

from retrieval import Retriever, build_prompt

# Fit a Retriever on 4–5 facts of your own, retrieve the top-2 for a question, and
# build_prompt them. Print the prompt and check the right facts were pulled in.

# Your implementation here:

Exercise 3: when does MMR help?

from retrieval import tfidf_matrix, encode_query, mmr

# Build a corpus with two near-duplicate relevant docs and one different relevant
# doc. Show that lambda=1.0 returns both duplicates but lambda=0.2 returns a
# diverse pair.

# Your implementation here:

Summary

Key takeaways:

  1. RAG conditions generation on retrieved text — encode → search → retrieve → augment → generate — so the model answers from a live corpus, not just frozen weights.
  2. Retrieval is nearest-neighbor search — encode documents and the query into one vector space and rank by cosine similarity.
  3. TF-IDF is the from-scratch encoder — term frequency times smoothed idf makes shared rare words count; dense embeddings are the same pipeline with a learned encoder.
  4. Chunk long documents into small overlapping passages so each retrieved unit is focused.
  5. MMR buys diversity — pick passages for relevance minus redundancy so the context isn’t three copies of one fact.
  6. RAG’s win is grounding — fresh, private, citable knowledge with no retraining, and a real dent in hallucination.

What’s Next

You can now ground a model in an external corpus. The natural next step is letting the model act: tool use and agents, where retrieval becomes one action among many (search, run code, call an API) inside a plan-act-observe loop — the same “put the right thing in the context” idea, now driven by the model itself.

Going Deeper

Core Papers:

Practical Resources: