Module 3 · RAG Done Right
RAG Done Right
Module 2 was about whether you can trust what the model says. This module is about whether the model was ever shown the right material. Retrieval is where most RAG systems quietly fail, and it's the one part you can measure without calling a model at all, so that's what you'll do: start from the simplest possible pipeline, score it, and improve it one change at a time, keeping only what the score justifies.
The labs share one answer key. A script generates fake product manuals as real PDFs, with tables, near-duplicate pages and a few image-only scans, and labels the exact page that answers each question. Because you know the right answer, every lab can report a number instead of an impression. Your numbers will differ from any quoted here, because they depend on the corpus the generator produces and the embedding model, so the checks tell you what to look for and ask you to record your own.
Concepts
Two halves, and the retrieval half decides everything
A RAG system (retrieval-augmented generation, the pattern named in Lewis et al.'s 2020 NeurIPS paper Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks) has an indexing half that runs once per document and a query half that runs on every question. The model only ever sees what retrieval hands it. If the right passage isn't in the top few results, no prompt can recover it, and the model will either say it can't find the answer or answer from the wrong passage.
That's why this module measures retrieval on its own, before any model is involved. Every lab scores recall@k: for each labeled question, did the correct page show up in the top k retrieved chunks? It's a plain pass/fail count you can compute in milliseconds and re-run after every change, which is exactly what makes it possible to improve a pipeline on evidence instead of vibes.
Parsing: most failures start before the embedding
A PDF isn't text. It's drawing instructions that happen to place glyphs on a page. A basic extractor like pypdf reads those glyphs back in drawing order, which is fine for a paragraph and ruinous for a table: every cell lands on its own line and the row structure disappears. A layout-aware parser (pymupdf4llm here; Docling, LlamaParse and Microsoft's markitdown are the same idea) detects the table and emits a real Markdown table. Lab 2 shows you the difference on your own screen.
A scanned PDF is worse: it's a stack of page images with no text layer at all, so a text extractor returns an empty string and the document is silently invisible to search. Scans need OCR (optical character recognition). Classic OCR engines like Tesseract work; so does sending a picture of the page to a vision-capable model and asking for a Markdown transcription, which also handles tables. Either way, check for empty pages explicitly. A pipeline that quietly indexes nothing looks exactly like one that works.
Chunking: the unit of retrieval
You can't embed a 200-page manual as one vector; the vector would be a blurry average of everything in it. So documents get cut into chunks, and each chunk is embedded and retrieved on its own. Cut too small and a chunk loses the surrounding meaning; cut too large and one vector has to represent several unrelated ideas, so none of them match a question well.
Three choices matter. Size (characters or tokens per chunk). Overlap, repeating the tail of one chunk at the start of the next so a fact straddling a boundary survives in at least one piece. And boundaries: splitting on the document's own structure (headings, paragraphs, whole tables) instead of every N characters. The strategy that works best depends on the document type, which is why Lab 3 compares them instead of recommending one.
Context: a chunk has to say where it came from
Take a specs table that lists pressure, cable length and part number but never names the product. Once it's cut out of its page, a human couldn't tell whose specs they are, and neither can an embedding model. The fix is cheap: prepend a short context line (document title, section, page) to every chunk before embedding it. Anthropic described a more elaborate variant of the same idea, having a model write each chunk's context, in its September 2024 post Introducing Contextual Retrieval. Lab 3 uses the simple, free version, built from the file name, and measures what it buys.
Vector indexes: HNSW, IVF and PQ
Finding the nearest vectors by comparing a question against every stored chunk is exact and simple, and it's fine up to tens of thousands of chunks. Beyond that you want an approximate nearest neighbor index, which trades a little accuracy for a lot of speed. Three names come up constantly:
HNSW (Malkov and Yashunin, Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs, IEEE TPAMI, 2018) links each vector to its neighbors in layered graphs and searches by hopping along links. It's fast and accurate and keeps the full vectors in memory plus the links. IVF (inverted file) clusters the vectors up front and, at query time, only scans the few clusters nearest the question; the nprobe setting is how many it checks. PQ (product quantization; Jégou, Douze and Schmid, IEEE TPAMI, 2011) compresses each vector into a short code, shrinking memory by an order of magnitude while making distances approximate. IVF and PQ are usually combined as IVF-PQ. Lab 4 builds both and measures the trade, using Meta's open-source Faiss library (Johnson, Douze and Jégou, Billion-scale similarity search with GPUs, IEEE Transactions on Big Data, 2019).
Hybrid search: meaning plus exact words
Embeddings match meaning, which is why "how long is the guarantee" can find a passage about the "warranty period". They are weaker on exact strings: part numbers, error codes, names, IDs. Keyword search with BM25 (the ranking function surveyed in Robertson and Zaragoza, The Probabilistic Relevance Framework: BM25 and Beyond, Foundations and Trends in Information Retrieval, 2009) is the opposite: blind to paraphrase, excellent on exact terms.
Hybrid search runs both and merges the two ranked lists. The merge that's become the default is reciprocal rank fusion (Cormack, Clarke and Büttcher, SIGIR 2009): each list gives every result a score of 1/(60 + rank), and you sum across lists. It only uses ranks, so you never have to reconcile BM25's score scale with cosine similarity's, and a chunk that both lists like rises to the top.
Question handling: split it, then filter it
Real questions are often two questions glued together ("compare the warranty on the Aster Probe and the Delta Sensor"). A single search for the whole sentence tends to find neither. The standard move is to have a model split a multi-part question into sub-questions, retrieve for each, and merge. Separately, many questions carry a metadata filter in plain sight: a product, a date range, a document type. Restricting the search to chunks whose metadata matches removes whole classes of near-identical distractors before ranking even starts. Lab 6 does the filter half with a simple rule; in production a model extracts the filter values from the question.
Reranking: a second, slower, smarter look
Embedding search scores a chunk and a question separately and then compares the vectors, which is fast but coarse. A cross-encoder reranker reads the question and one chunk together and outputs a relevance score, which is far more precise and far too slow to run over every chunk. So the standard shape is two stages: cheap retrieval pulls a pool (top 20), the reranker re-orders just that pool, and you keep the top 5. You pay latency on every query, so Lab 6 measures both what you gain and what it costs.
Hands-on Labs
- Install this module's dependencies. They go in the same shared environment as the earlier modules (with
uv, useuv addinstead ofpip install).pip install pypdf pymupdf4llm pypdfium2 reportlab numpy rank-bm25 faiss-cpu sentence-transformers - Generate a corpus you can grade. This script writes fake product manuals as real PDFs (each has a warranty paragraph, a specs table, and a maintenance page), plus
rag/questions.json: labeled questions, each with the exact document and page that answers it and the expected answer. Because we placed every fact ourselves, the answer key is correct by construction. Every manual reuses the same filler sentences on purpose, so lots of pages look alike, like a real set of sibling documents. Every 20th manual is saved as an image-only scan with no text layer."""Make a corpus of fake product manuals plus a labeled question set. Every fact sits on a known page, so the answer key is exact by construction.""" import json, random, sys from pathlib import Path import pypdfium2 as pdfium from reportlab.lib import colors from reportlab.lib.pagesizes import letter from reportlab.lib.styles import getSampleStyleSheet from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak N_DOCS = int(sys.argv[1]) if len(sys.argv) > 1 else 12 OUT = Path("rag/corpus"); OUT.mkdir(parents=True, exist_ok=True) random.seed(7) ADJ = ["Aster", "Birch", "Cobalt", "Delta", "Ember", "Fjord", "Garnet", "Harbor", "Ivory", "Juniper"] NOUN = ["Gauge", "Probe", "Meter", "Sensor", "Valve", "Pump", "Relay", "Logger", "Scale", "Analyzer"] FILLER = [ "Read all safety notices before installing or operating this unit.", "Use only genuine Larkfield accessories; third-party parts void the warranty.", "Store the unit in a dry location between {lo} and {hi} degrees Celsius.", "Clean the housing with a damp cloth. Never immerse the unit in liquid.", "Firmware updates are published quarterly on the Larkfield support portal.", "Contact an authorized service center for any repair not covered in this manual.", "Dispose of the unit according to local electronic waste regulations.", "Keep this manual with the unit for the life of the product.", ] def filler(n): return [random.choice(FILLER).format(lo=random.randint(-20, -5), hi=random.randint(40, 70)) for _ in range(n)] names = [f"{a} {n}" for a in ADJ for n in NOUN] random.shuffle(names) questions = [] for i in range(N_DOCS): name = names[i] code = f"LK-{random.randint(1000, 9999)}-{random.choice('ABCDEFGH')}" warranty = random.choice([12, 18, 24, 36, 48]) pressure = random.randint(40, 400) cable = random.randint(2, 30) interval = random.choice([3, 6, 9, 12, 24]) doc_id = f"{i:03d}_{name.lower().replace(' ', '_')}" ss = getSampleStyleSheet() story = [Paragraph(f"Larkfield {name} User Manual", ss["Title"]), Paragraph(f"Part number {code}", ss["Heading3"])] story += [Paragraph(s, ss["BodyText"]) for s in filler(16)] story += [Paragraph(f"The Larkfield {name} is covered by a limited warranty of {warranty} months from the date of purchase.", ss["BodyText"]), PageBreak()] story += [Paragraph("Technical specifications", ss["Heading2"])] t = Table([["Parameter", "Value", "Unit"], ["Maximum pressure", str(pressure), "bar"], ["Cable length", str(cable), "m"], ["Part number", code, ""]]) t.setStyle(TableStyle([("GRID", (0, 0), (-1, -1), 0.5, colors.grey), ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey)])) story += [t, Spacer(1, 12)] + [Paragraph(s, ss["BodyText"]) for s in filler(18)] + [PageBreak()] story += [Paragraph("Maintenance", ss["Heading2"])] + [Paragraph(s, ss["BodyText"]) for s in filler(12)] story += [Paragraph(f"Recalibrate the {name} every {interval} months to keep readings within tolerance.", ss["BodyText"])] story += [Paragraph(s, ss["BodyText"]) for s in filler(12)] pdf_path = OUT / f"{doc_id}.pdf" SimpleDocTemplate(str(pdf_path), pagesize=letter).build(story) questions += [ {"q": f"How long is the warranty on the Larkfield {name}?", "doc": pdf_path.name, "pages": [1], "answer": f"{warranty} months"}, {"q": f"What is the maximum pressure of the {name}?", "doc": pdf_path.name, "pages": [2], "answer": f"{pressure} bar"}, {"q": f"How often should I recalibrate the {name}?", "doc": pdf_path.name, "pages": [3], "answer": f"{interval} months"}, ] if i % 4 == 0: # an exact-identifier question: a part number in a table cell questions.append({"q": f"Which product is part number {code}?", "doc": pdf_path.name, "pages": [1, 2], "answer": name}) # printed on both pages if i % 20 == 3: # make a scanned copy: page images only, no text layer src = pdfium.PdfDocument(str(pdf_path)) imgs = [src[p].render(scale=1.5).to_pil().convert("L") for p in range(len(src))] imgs[0].save(OUT / f"{doc_id}_scan.pdf", save_all=True, append_images=imgs[1:]) pdf_path.unlink() for q in questions: if q["doc"] == pdf_path.name: q["doc"] = f"{doc_id}_scan.pdf" Path("rag/questions.json").write_text(json.dumps(questions, indent=1)) print(f"{N_DOCS} docs, {len(questions)} labeled questions") - Run it. The default is 12 manuals; the project at the end uses 100.
python rag/generate_corpus.py 12 - Create
rag/common.py, the helpers every later lab imports: a page loader, an embedder that caches every vector on disk (so re-running a lab doesn't re-bill you), cosine search, and therecall_at_kscorer.import hashlib, json, os from pathlib import Path import numpy as np import pypdf from dotenv import load_dotenv load_dotenv() CACHE = Path("rag/.embed_cache"); CACHE.mkdir(exist_ok=True) def load_pages(corpus="rag/corpus"): """Naive parser: one (doc, page_number, text) per PDF page.""" pages = [] for pdf in sorted(Path(corpus).glob("*.pdf")): for n, page in enumerate(pypdf.PdfReader(pdf).pages, start=1): pages.append({"doc": pdf.name, "page": n, "text": page.extract_text() or ""}) return pages def embed(texts, model="text-embedding-3-small"): """Embed a list of strings, caching each one on disk so re-runs cost nothing.""" from openai import OpenAI client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) out, todo = {}, [] for t in texts: f = CACHE / (hashlib.sha256((model + t).encode()).hexdigest() + ".json") if f.exists(): out[t] = json.loads(f.read_text()) else: todo.append((t, f)) for i in range(0, len(todo), 100): batch = todo[i:i + 100] resp = client.embeddings.create(model=model, input=[t for t, _ in batch]) for (t, f), d in zip(batch, resp.data): f.write_text(json.dumps(d.embedding)); out[t] = d.embedding return np.array([out[t] for t in texts], dtype="float32") def top_k_vector(q_vec, chunk_vecs, k): """Cosine similarity via normalized dot product; returns chunk indexes, best first.""" a = chunk_vecs / np.linalg.norm(chunk_vecs, axis=1, keepdims=True) b = q_vec / np.linalg.norm(q_vec) return list(np.argsort(-(a @ b))[:k]) def recall_at_k(questions, retrieve, chunks, k=5): """retrieve(question, k) -> list of chunk indexes. A hit means any of the top k chunks comes from the labeled doc AND one of its labeled pages.""" hits, misses = 0, [] for q in questions: got = [(chunks[i]["doc"], chunks[i]["page"]) for i in retrieve(q["q"], k)] if any((q["doc"], pg) in got for pg in q["pages"]): hits += 1 else: misses.append(q["q"]) return hits / len(questions), misses - Create
rag/naive_rag.py: the simplest thing that could work. Plainpypdftext, fixed 500-character chunks, one embedding per chunk, top-5 by cosine similarity. It scores recall first, then answers one question with citations so you can see the whole loop.import json from common import load_pages, embed, top_k_vector, recall_at_k from openai import OpenAI import os def chunk_fixed(pages, size=500): chunks = [] for p in pages: for start in range(0, len(p["text"]), size): piece = p["text"][start:start + size].strip() if piece: chunks.append({"doc": p["doc"], "page": p["page"], "text": piece}) return chunks pages = load_pages() chunks = chunk_fixed(pages) chunk_vecs = embed([c["text"] for c in chunks]) questions = json.load(open("rag/questions.json")) def retrieve(question, k): return top_k_vector(embed([question])[0], chunk_vecs, k) def answer(question, k=5): ctx = [chunks[i] for i in retrieve(question, k)] context = "\n\n".join(f"[{c['doc']} p.{c['page']}]\n{c['text']}" for c in ctx) client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) resp = client.chat.completions.create( model="gpt-4.1-mini", max_tokens=200, messages=[{"role": "system", "content": "Answer using only the context. Cite the [doc p.N] " "tag for each fact. If the context does not contain the answer, say so."}, {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}]) return resp.choices[0].message.content if __name__ == "__main__": score, misses = recall_at_k(questions, retrieve, chunks, k=5) print(f"{len(chunks)} chunks | recall@5 = {score:.0%} ({len(questions)} questions)") for m in misses[:8]: print(" missed:", m) print("\n", answer(questions[0]["q"])) - Run it and write down the recall@5 number and the missed questions. This is your baseline; every later lab has to beat it.
python rag/naive_rag.py
- Compare the two parsers side by side on a spec page. Pass any non-scanned PDF from
rag/corpus/.import sys import pypdf, pymupdf4llm path = sys.argv[1] # e.g. rag/corpus/000_delta_sensor.pdf print("=== pypdf, page 2 ===") print(pypdf.PdfReader(path).pages[1].extract_text()[:300]) print("\n=== pymupdf4llm, page 2 ===") print(pymupdf4llm.to_markdown(path, pages=[1], use_ocr=False)[:300])python rag/parse_compare.py rag/corpus/000_delta_sensor.pdf - Read both outputs. The
pypdfversion puts every cell on its own line, so "Maximum pressure", "282" and "bar" are three separate lines with nothing marking them as one row. Thepymupdf4llmversion gives you a Markdown table where the row survives. - Now the scan. Confirm a scanned manual really has no text for a basic extractor, then see what the layout-aware path does with it.
That prints an empty string. Your naive pipeline from Lab 1 indexed this document as nothing.
python -c "import pypdf,glob; f=glob.glob('rag/corpus/*_scan.pdf')[0]; print(f, repr(pypdf.PdfReader(f).pages[0].extract_text()))" - Create
rag/parsing.py. It usespymupdf4llmfor every page and, for any page that comes back empty, renders the page to an image and asks a vision model to transcribe it to Markdown. Transcriptions are cached, so each scanned page is paid for once. (If you'd rather use a local OCR engine like Tesseract, swap it in at the same spot; the contract is "page in, Markdown out".)import base64, hashlib, os from pathlib import Path import pymupdf4llm import pypdfium2 as pdfium from dotenv import load_dotenv load_dotenv() OCR_CACHE = Path("rag/.ocr_cache"); OCR_CACHE.mkdir(exist_ok=True) def ocr_page(pdf_path, page_index): """Scanned page -> Markdown, by sending a picture of the page to a vision model. Cached on disk so each page is transcribed (and paid for) once.""" f = OCR_CACHE / f"{Path(pdf_path).stem}_p{page_index + 1}.md" if f.exists(): return f.read_text() from openai import OpenAI png = pdfium.PdfDocument(str(pdf_path))[page_index].render(scale=2).to_pil() png.save(f.with_suffix(".png")) b64 = base64.b64encode(f.with_suffix(".png").read_bytes()).decode() resp = OpenAI(api_key=os.environ["OPENAI_API_KEY"]).chat.completions.create( model="gpt-4.1-mini", max_tokens=1500, messages=[{"role": "user", "content": [ {"type": "text", "text": "Transcribe this page to Markdown. Keep headings, and keep tables as " "Markdown tables. Output only the transcription."}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}]}]) f.write_text(resp.choices[0].message.content) return f.read_text() def load_pages_md(corpus="rag/corpus"): """Layout-aware parse: one Markdown string per page. Tables stay tables; pages with no text layer (scans) fall back to OCR.""" pages = [] for pdf in sorted(Path(corpus).glob("*.pdf")): parsed = pymupdf4llm.to_markdown(str(pdf), page_chunks=True, use_ocr=False) for n, md in enumerate(parsed, start=1): text = md["text"] if not text.strip(): text = ocr_page(pdf, n - 1) pages.append({"doc": pdf.name, "page": n, "text": text}) return pages - First save a copy of the Lab 1 version (
cp rag/naive_rag.py rag/naive_rag_v1.py), since the project needs that baseline later. Then changerag/naive_rag.pyto build its chunks fromload_pages_md()instead ofload_pages()(keep the fixed 500-character chunker) and re-run it. The only thing that changed is the parser.
rag/.ocr_cache/ and read the transcription: a wrong or empty transcription there is an OCR problem, not a retrieval problem, and it's worth knowing which one you have. Note how much of the improvement came from parsing alone, with no change to embeddings or search. Also remember the licensing note from Module 2: pymupdf4llm is AGPL-3.0 or a paid commercial license.- Create
rag/chunking.pywith three chunkers and the context-prefix helper: fixed-size with an optional overlap; a structure-aware chunker that splits on blank lines (so a Markdown table stays whole) and packs neighbors up to a size cap; andadd_context, which prepends "From the Larkfield <product> manual, page N:" to every chunk.import re def chunk_fixed(pages, size=500, overlap=0): out, step = [], size - overlap for p in pages: for start in range(0, len(p["text"]), step): piece = p["text"][start:start + size].strip() if piece: out.append({"doc": p["doc"], "page": p["page"], "text": piece}) return out def chunk_structure(pages, max_chars=800): """Split on blank lines (paragraphs, headings, whole tables), then pack neighbors together until max_chars. A table never gets cut in half unless it alone exceeds the cap.""" out = [] for p in pages: blocks = [b.strip() for b in re.split(r"\n\s*\n", p["text"]) if b.strip()] buf = "" for b in blocks: if buf and len(buf) + len(b) > max_chars: out.append({"doc": p["doc"], "page": p["page"], "text": buf}); buf = "" buf = f"{buf}\n\n{b}" if buf else b if buf: out.append({"doc": p["doc"], "page": p["page"], "text": buf}) return out def add_context(chunks): """Prefix every chunk with where it came from. A specs table that never mentions the product name becomes findable once the prefix says whose specs they are.""" out = [] for c in chunks: product = re.sub(r"^\d+_|_scan|\.pdf", "", c["doc"]).replace("_", " ").title() out.append({**c, "text": f"From the Larkfield {product} manual, page {c['page']}:\n{c['text']}"}) return out - Create
rag/chunk_experiment.py. It scores four strategies on the same questions and the same parse, so the chunking is the only thing that changes.import json from common import embed, top_k_vector, recall_at_k from parsing import load_pages_md from chunking import chunk_fixed, chunk_structure, add_context pages = load_pages_md() questions = json.load(open("rag/questions.json")) strategies = { "fixed 500, no overlap": chunk_fixed(pages, 500, 0), "fixed 500, overlap 100": chunk_fixed(pages, 500, 100), "structure-aware, cap 800": chunk_structure(pages, 800), "structure + context prefix": add_context(chunk_structure(pages, 800)), } for name, chunks in strategies.items(): vecs = embed([c["text"] for c in chunks]) retrieve = lambda q, k: top_k_vector(embed([q])[0], vecs, k) score, _ = recall_at_k(questions, retrieve, chunks, k=5) avg = sum(len(c["text"]) for c in chunks) / len(chunks) print(f"{name:28s} chunks={len(chunks):4d} avg_chars={avg:5.0f} recall@5={score:.0%}") - Run it. The first run embeds four different chunk sets; the disk cache means re-runs are free.
python rag/chunk_experiment.py - For each strategy, pick two questions it still misses and read the retrieved chunks next to the correct page. Ask what the chunk was missing: context, a table row cut in half, or a boundary that split a sentence.
- Create
rag/index_types.py. It builds 100,000 synthetic vectors (clustered, because real embeddings clump by topic and uniform random vectors would make every index look worse than it will on documents), computes the exact nearest neighbors as the answer key, then builds HNSW and IVF-PQ indexes and measures build time, search speed and recall against that key.import time, numpy as np, faiss D, N, NQ, K = 128, 100_000, 200, 10 rng = np.random.default_rng(0) # Clustered data, not uniform noise: real embeddings clump by topic, and uniform # random vectors make every index look worse than it will on real documents. centers = rng.normal(size=(200, D)).astype("float32") data = (centers[rng.integers(0, 200, N)] + 0.35 * rng.normal(size=(N, D))).astype("float32") queries = (centers[rng.integers(0, 200, NQ)] + 0.35 * rng.normal(size=(NQ, D))).astype("float32") faiss.normalize_L2(data); faiss.normalize_L2(queries) def bench(name, index, train=False, setup=None): t0 = time.time() if train: index.train(data) index.add(data) build = time.time() - t0 if setup: setup(index) t0 = time.time() _, ids = index.search(queries, K) qps = NQ / (time.time() - t0) recall = np.mean([len(set(ids[i]) & set(truth[i])) / K for i in range(NQ)]) print(f"{name:10s} build={build:5.1f}s qps={qps:8.0f} recall@{K}={recall:.2f}") flat = faiss.IndexFlatIP(D); flat.add(data) # exact search = the answer key _, truth = flat.search(queries, K) print(f"raw vectors: {N * D * 4 / 1e6:.0f} MB") bench("flat", faiss.IndexFlatIP(D)) bench("HNSW", faiss.IndexHNSWFlat(D, 32, faiss.METRIC_INNER_PRODUCT), setup=lambda ix: setattr(ix.hnsw, "efSearch", 64)) ivfpq = faiss.IndexIVFPQ(faiss.IndexFlatIP(D), D, 256, 32, 8, faiss.METRIC_INNER_PRODUCT) bench("IVF-PQ", ivfpq, train=True, setup=lambda ix: setattr(ix, "nprobe", 32)) print(f"IVF-PQ stores about {N * 32 / 1e6:.1f} MB of codes (32 bytes per vector) vs {N * D * 4 / 1e6:.0f} MB raw") - Run it.
python rag/index_types.py - Now turn the knobs, one at a time, and re-run: raise HNSW's
efSearchfrom 64 to 256; raise IVF-PQ'snprobefrom 32 to 128; change PQ's code size from 32 to 64 bytes per vector. For each, write down what moved: recall, speed, or memory.
0.97; IVF-PQ stored the same vectors in about 3 MB but recalled only about 0.30. Speeds depend entirely on your machine, so compare them to each other rather than to ours. The point of the lab is the shape of the trade: HNSW buys accuracy and speed with memory; PQ buys memory with accuracy, which you win back with a higher nprobe, bigger codes, or by re-scoring PQ's shortlist against the full vectors. At this module's corpus size (a few hundred chunks), none of this is needed, and an exact search is the right answer. You're building the instinct for when the corpus is a million chunks.- Create
rag/hybrid.py. It builds a BM25 index next to the vector index over the same chunks (structure-aware, with context prefix, from Lab 3), definesrrfto merge two ranked lists, and prints recall for vector-only, keyword-only and hybrid, followed by the specific questions hybrid rescued.import json, re from rank_bm25 import BM25Okapi from common import embed, top_k_vector, recall_at_k from parsing import load_pages_md from chunking import chunk_structure, add_context def tokenize(text): return re.findall(r"\w+", text.lower()) pages = load_pages_md() chunks = add_context(chunk_structure(pages, 800)) vecs = embed([c["text"] for c in chunks]) bm25 = BM25Okapi([tokenize(c["text"]) for c in chunks]) questions = json.load(open("rag/questions.json")) def vector_search(q, k): return list(top_k_vector(embed([q])[0], vecs, k)) def keyword_search(q, k): scores = bm25.get_scores(tokenize(q)) return sorted(range(len(chunks)), key=lambda i: -scores[i])[:k] def rrf(rankings, k=60): """Reciprocal rank fusion: each list votes 1/(k+rank) for every chunk it returns.""" fused = {} for ranking in rankings: for rank, i in enumerate(ranking, start=1): fused[i] = fused.get(i, 0) + 1 / (k + rank) return sorted(fused, key=fused.get, reverse=True) def hybrid_search(q, k, pool=20): return rrf([vector_search(q, pool), keyword_search(q, pool)])[:k] if __name__ == "__main__": for name, fn in [("vector", vector_search), ("keyword", keyword_search), ("hybrid", hybrid_search)]: score, misses = recall_at_k(questions, fn, chunks, k=5) print(f"{name:8s} recall@5 = {score:.0%}") # Which questions does fusion rescue? Compare per-question hits. def hit(fn, q): got = [(chunks[i]["doc"], chunks[i]["page"]) for i in fn(q["q"], 5)] return any((q["doc"], pg) in got for pg in q["pages"]) print("\nFixed by hybrid (vector missed, hybrid found):") for q in questions: if not hit(vector_search, q) and hit(hybrid_search, q): print(" ", q["q"]) - Run it.
python rag/hybrid.py - Look at the "Fixed by hybrid" list. Is it the exact-identifier questions ("Which product is part number LK-...?")? If it's something else, work out why keyword search rescued it. Then check the reverse: are there questions vector search found that keyword search missed? Those are the paraphrases hybrid exists to protect.
- Try the one fusion knob. Change the
k=60constant inrrfto 10 and then 200 and watch whether anything moves. (It's a smoothing constant from the original paper; on a small corpus it rarely matters, and knowing that saves you from tuning it.)
- Create
rag/rerank_filter.py. It takes the hybrid pool (top 20), optionally filters it to the product named in the question, then reorders it with a small open cross-encoder,cross-encoder/ms-marco-MiniLM-L-6-v2, and keeps the top 5. The product filter here is a deliberately simple rule (match the product's name in the question). In production you'd have a model extract it.import json, re, time from sentence_transformers import CrossEncoder from common import recall_at_k from hybrid import chunks, questions, hybrid_search reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2") def product_slug(question): """Stand-in for 'extract the product from the question'. Matches a doc whose name (e.g. 'delta_sensor') appears in the question text.""" q = re.sub(r"[^a-z ]", "", question.lower()) for doc in {c["doc"] for c in chunks}: name = re.sub(r"^\d+_|_scan|\.pdf", "", doc).replace("_", " ") if name in q: return doc return None def retrieve_rerank(q, k, pool=20, use_filter=False): candidates = hybrid_search(q, pool) if use_filter and (doc := product_slug(q)): candidates = [i for i in candidates if chunks[i]["doc"] == doc] or candidates scores = reranker.predict([(q, chunks[i]["text"]) for i in candidates]) return [i for _, i in sorted(zip(scores, candidates), reverse=True)][:k] if __name__ == "__main__": for label, fn in [("hybrid only", hybrid_search), ("hybrid + rerank", retrieve_rerank), ("hybrid + rerank + filter", lambda q, k: retrieve_rerank(q, k, use_filter=True))]: t0 = time.time() score, _ = recall_at_k(questions, fn, chunks, k=5) ms = (time.time() - t0) / len(questions) * 1000 print(f"{label:26s} recall@5={score:.0%} {ms:.0f} ms/question") - Run it. The first run downloads the reranker model (tens of megabytes) and loads PyTorch, so expect a pause.
python rag/rerank_filter.py - Record both columns for each row: recall@5 and milliseconds per question. The reranker adds latency to every query, and whether the recall gain is worth that is a product decision, not a technical one.
- Break the filter on purpose: ask a question that names no product ("What is the warranty period?"). Confirm the code falls back to the unfiltered pool instead of returning nothing. A filter that can silently empty your results is a bug waiting for a user.
Project
The labs ran on 12 manuals. This project scales the winning pipeline to 100, wraps it in a function a chat app can call, answers with citations that point at a specific page of the source PDF, caches repeat questions, and writes down what each upgrade was worth.
- Generate the bigger corpus and re-run your saved Lab 1 baseline (
naive_rag_v1.py) on it first. The number you get here is the "before" in your results table.Parsing the five scans costs a few vision-model calls per page, once; they are cached after that.python rag/generate_corpus.py 100 python rag/naive_rag_v1.py - Create
rag/pipeline.py, the one function the app calls. It checks a semantic cache first (the same idea as Module 1, Lab 6: embed the question, return the stored answer if the nearest earlier question clears your threshold), otherwise runs hybrid retrieval, product filter and rerank, then asks the model for an answer that cites[doc p.N]tags. It reports latency and token counts for every call.Use the cache threshold you found safe in Module 1, not a number copied from here. The cache must hold the sources alongside the answer, so a cache hit still shows citations.import os, time import numpy as np from openai import OpenAI from common import embed from hybrid import chunks from rerank_filter import retrieve_rerank client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) CACHE_THRESHOLD = 0.95 # Re-tune this on YOUR questions. Two questions that differ only # in the product name score very high, and must NOT share an answer. _cache = [] # list of (question_vector, result_dict) def _cache_lookup(q_vec): for vec, result in _cache: sim = float(vec @ q_vec / (np.linalg.norm(vec) * np.linalg.norm(q_vec))) if sim >= CACHE_THRESHOLD: return result return None def ask(question): t0 = time.perf_counter() q_vec = embed([question])[0] hit = _cache_lookup(q_vec) if hit: return {**hit, "cached": True, "latency_ms": (time.perf_counter() - t0) * 1000, "input_tokens": 0, "output_tokens": 0} top = [chunks[i] for i in retrieve_rerank(question, 5, use_filter=True)] context = "\n\n".join(f"[{c['doc']} p.{c['page']}]\n{c['text']}" for c in top) resp = client.chat.completions.create( model="gpt-4.1-mini", max_tokens=300, messages=[{"role": "system", "content": "Answer using only the context. After each fact, cite its " "source tag exactly as written, like [doc.pdf p.2]. If the context does not contain " "the answer, say so and do not guess."}, {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}]) result = {"answer": resp.choices[0].message.content, "sources": [{"doc": c["doc"], "page": c["page"]} for c in top], "contexts": [c["text"] for c in top]} _cache.append((q_vec, result)) return {**result, "cached": False, "latency_ms": (time.perf_counter() - t0) * 1000, "input_tokens": resp.usage.prompt_tokens, "output_tokens": resp.usage.completion_tokens} if __name__ == "__main__": for q in ["How long is the warranty on the Larkfield Delta Sensor?", "How long is the warranty on the Larkfield Delta Sensor?", # repeat -> cache hit "How long is the warranty on the Larkfield Cobalt Pump?"]: # similar but different -> miss r = ask(q) print(f"cached={r['cached']!s:5} {r['latency_ms']:6.0f} ms in={r['input_tokens']} out={r['output_tokens']} {q}") print(" ", r["answer"][:120].replace("\n", " ")) - Build your "before and after" table. Re-run the recall measurement after each upgrade in this order, and add a row each time: naive baseline, layout parser plus OCR, structure-aware chunks, context prefix, hybrid search, rerank, filter. Save it as
rag/results.md. The goal isn't a high final number; it's knowing which single change was worth the most, and which ones barely mattered. - Put a thin app on top. Any web framework works. The one requirement that makes citations real: serve the PDFs, and make each citation a link of the form
/corpus/<doc>#page=<N>, which browsers' built-in PDF viewers jump straight to.
Done when:
- Clicking a citation opens the correct PDF at the correct page, for at least ten questions you spot-check by hand, including one from a scanned manual.
- Recall@5 on the 100-manual corpus beats your naive baseline, and
rag/results.mdnames the single upgrade that gained the most. - A repeated question is answered from the cache (visibly faster, no model call), and a same-topic-but-different question is not.
- Cost and latency per question are reported for both a cache miss and a cache hit.
Sources
- Lewis, P., et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS 2020.
- Robertson, S., & Zaragoza, H. (2009). The Probabilistic Relevance Framework: BM25 and Beyond. Foundations and Trends in Information Retrieval, 3(4).
- Cormack, G. V., Clarke, C. L. A., & Büttcher, S. (2009). Reciprocal Rank Fusion Outperforms Condorcet and Individual Rank Learning Methods. SIGIR 2009.
- Malkov, Y. A., & Yashunin, D. A. (2018). Efficient and Robust Approximate Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs. IEEE TPAMI, 42(4).
- Jégou, H., Douze, M., & Schmid, C. (2011). Product Quantization for Nearest Neighbor Search. IEEE TPAMI, 33(1).
- Johnson, J., Douze, M., & Jégou, H. (2019). Billion-scale Similarity Search with GPUs. IEEE Transactions on Big Data, 7(3).
- Anthropic (2024). Introducing Contextual Retrieval. anthropic.com/news/contextual-retrieval.