Retrieval lab
Retrieval lab
You build retrieval_lab, a small hybrid retrieval pipeline over Postgres: a Markdown-aware
chunker, a deterministic hashing embedder, a pgvector table with a generated full-text column,
vector and full-text search fused with reciprocal rank fusion (RRF), a reranker with MMR,
retrieval metrics on a golden set, and an answer assembler that cites its sources. The last stages
swap in local models (fastembed) and ask Claude for cited answers.
The corpus is fixtures/corpus: 29 Markdown documents (a snapshot of course lessons and interview-angle
pages) and three public-domain stories, with golden.yaml, 48 queries of three kinds
(factual, conceptual, lexical). A chunk is relevant to a golden query when it belongs to the
query’s doc_id and its text contains the query’s anchor, case-insensitive and
whitespace-normalised. See fixtures/corpus/README.md.
retrieval_lab/types.py is given and shared by every stage (Document, Chunk, Hit, Ranked,
GoldenQuery, Metrics, Draft, Faithfulness, Answer). Write the stage modules against it;
do not change it.
How to run
Run every command below in a terminal at the root of your clone of this repository, after a one-time uv sync there; uv run miubais doctor lists what this machine has.
uv run miubais start m4-retrieval-labThis copies course/m4-retrieval-lab/starter/ to workspace/m4-retrieval-lab/ at the repo root. That copy is yours: the stubs under workspace/m4-retrieval-lab/src/retrieval_lab/ raise NotImplementedError until you fill them in, and the judgment stage asks for a file in the workspace root, named in its description below. The tests stay in course/m4-retrieval-lab/tests/ and import your package from the workspace. start --force replaces the workspace with a fresh copy.
uv run miubais check m4-retrieval-labRuns the stages in order and stops at the first failure, printing its pytest output, so stage 02 waits for stage 01. --stage <stage id> runs one stage and --all keeps going past a failure. A stage whose requires (in stages.yaml) are missing on your machine is skipped, not failed. Results are kept in ~/.miubais/progress.json; uv run miubais reset m4-retrieval-lab forgets them (--all for every exercise), and start --force does the same along with the fresh copy.
uv run miubais hint m4-retrieval-labTakes the last stage check saw fail, its test output and your diff against the starter, and asks the tutor on your Claude Code login for a nudge of under 90 words that is checked not to hand you the solution. If the model cannot be reached, you get a pointer to the failing test instead.
The Run check and Hint buttons on the site’s exercise page do the same while uv run miubais serve is running, and share the progress file with the terminal.
Stages
-
01-chunker:chunker.load_documents(corpus_dir)returns oneDocument(doc_id, title, text)per file ofdocs/*.mdthenstories/*.txt, each sorted by name;doc_idis the file stem andtitlethe first H1 (or the first non-blank line of a story).chunker.chunk(doc, *, max_chars=1200, overlap=1)follows these rules:- A Markdown heading (
#to######followed by a space, outside a fence) starts a new section; the heading line is the first block of its section andheading_pathjoins the open headings with>(Title > Section A > Sub A.1; text before the first heading has""). A chunk never spans two sections. - Inside a section, blocks are separated by blank lines. A fenced code block (a line starting
with three or more backticks or tildes, up to the closing line of the same character) is one
block, blank lines and
#lines included; an unclosed fence runs to the end of the document. - Blocks are packed greedily: a chunk grows while the slice from its first block’s start to the
next block’s end stays within
max_chars. A block longer thanmax_chars(a long fence, a long list) becomes one oversized chunk; it is never split. - The next chunk starts
overlapblocks before the previous one ended, but only when the carried blocks still leave room for at least one new block; otherwise it starts with the new block. start_char/end_charare offsets intodoc.textanddoc.text[start_char:end_char] == text;chunk_idisf"{doc_id}#{index:03d}", numbered from 0 across the document.
The tests check heading paths, packing with and without overlap on invented documents, that a 2,000-character fence is one chunk, offsets, ids and determinism on the corpus, that every non-blank character is covered, that no chunk splits a fence, and that every golden anchor lies inside at least one chunk of its document (so recall can reach 1.0).
- A Markdown heading (
-
02-embed-and-store(Docker):embed.Embedderis a protocol withdims,nameandembed(texts) -> np.ndarrayof shape(len(texts), dims).embed.HashEmbedder(dims=256)is pinned exactly: tokens are the lowercase words ofre.findall(r"\w+", text.lower()), then the word bigramsf"{a} {b}", then every three-character slice of every word; each token is hashed withhashlib.blake2b(token.encode(), digest_size=8), read as a big-endian integerh, and adds+1at indexh % dims(or-1when the top bith >> 63is set); the vector is L2-normalised (float32). It must give the same vector in every process, so no built-inhash().embed.FastEmbedder(model="BAAI/bge-small-en-v1.5", cache_dir=None)hasdims384 and anamethat includes the model; it importsfastembedonly when it embeds (written in stage 08).store.Schema.create(conn, dims)creates, if missing, thevectorextension andCREATE TABLE chunks (chunk_id text PRIMARY KEY, doc_id text, heading_path text,start_char int, end_char int, text text,tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED,embedding vector(<dims>));CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops);CREATE INDEX ON chunks USING gin (tsv);and raises
ValueErrorwhenchunksalready exists with a different dimension.store.upsert(conn, chunks, embedder, batch=64)embeds and writes in batches withINSERT ... ON CONFLICT (chunk_id) DO UPDATE, so a second run changes nothing and an edited chunk is replaced; it raisesValueErrorbefore writing whenembedder.dimsdiffers from the table.store.count(conn)returns the row count. The tests check the pinned recipe against an independent implementation, determinism across processes with differentPYTHONHASHSEED, unit norm, the columns and both indexes, idempotent upserts, the stored vectors, and both mismatches. -
03-hybrid-retrieval(Docker): every search returnsHit(chunk_id, score, rank)with ranks fromhybrid.vector_search(conn, embedder, query, n=50)orders by cosine distance<=>and scores1 - distance, and returns[]for a query without words (a zero vector has no cosine); an HNSW scan returns at mosthnsw.ef_searchrows (default 40), so raise it tonfor the query (SET LOCALorset_config(..., true)inside a transaction).hybrid.text_search(conn, query, n=50)parses the query withwebsearch_to_tsquery('english', query), which never raises on user input, then splits its text form at the top-level&and|: the positive terms (words and quoted phrases) are ORed into one queryq, and each negated term (-word,-"a phrase", printed as!...) stays an exclusion,tsv @@ q AND NOT tsv @@ (neg1 | neg2 ...).websearch_to_tsqueryANDs every word, and a question in natural language rarely has all its words in one chunk; turning every&into|would instead turn!terminto “any chunk without the term”. A query with no positive term returns[]. It ranks withts_rank_cd(tsv, q, 32), ties bychunk_id.hybrid.rrf(lists, k=60)fuses by rank,score(d) = Σ 1 / (k + rank_r(d))over the lists that containd, sorted by score thenchunk_id, and returnsRanked(chunk_id, score, rank, ranks)withranksholding the rank in each input list (Nonewhen absent).hybrid.search(conn, embedder, query, n=50, top_k=10)fuses the vector and full-text top-nand returns the firsttop_k. The tests check cosine scores against the embedder, full-text ranks, that queries with operators and quotes ("exact phrase" -term OR x, a stray',&|!<->) never raise and still return vector results, that a query without words ("",-) returns[]everywhere, that-termand-"a phrase"exclude every chunk containing them while still returning results, the RRF formula and tie-break, that a chunk first in one list beats a chunk 30th in the other list however large that list’s raw score, that a chunk 30th in both lists (2/90) outranks one first in a single list (1/61), that the lexical golden queryq40is found by full text and missed by vectors, and that fused recall@10 on the golden set is at least that of each side and above one of them.
-
04-reranker(Docker):rerank.Rerankeris a protocol withrerank(query, chunks, top_k) -> list[Ranked], best first, ranks from 1, ties kept in input order.rerank.LexicalReranker(k1=1.2, b=0.75)scores BM25 over the candidate set it is given: query terms are the distinct lowercase words ofre.findall(r"\w+", query.lower())(chunk text is tokenised the same way); for a termtwith document frequencydfamong theNcandidates,idf = ln(1 + (N - df + 0.5) / (df + 0.5)), and a chunk ofLwords (averageavgover the candidates) scoresΣ idf · tf · (k1 + 1) / (tf + k1 · (1 - b + b · L / avg))over the query terms it contains.rerank.CrossEncoderReranker(model="Xenova/ms-marco-MiniLM-L-6-v2", cache_dir=None)wraps fastembed’sTextCrossEncoderand imports fastembed only when it ranks a non-empty list (used in stage 08).rerank.mmr(query_vec, candidates, embeddings, *, lam=0.7, top_k)picks greedily the candidate with the highestlam · relevance - (1 - lam) · max cosine to the chunks already picked(the second term is 0 for the first pick; ties go to the earlier candidate). Relevance is the reranker score scaled into[0, 1]:score / maxwhen no score is negative, min-max otherwise, and the cosine toquery_vecwhen every score is equal. It keeps each candidate’s reranker score and renumbers the ranks.rerank.fetch(conn, chunk_ids)returns{chunk_id: (Chunk, embedding)}from the table, andrerank.pipeline(conn, embedder, reranker, query, *, n=50, top_k=10, mmr_lambda=0.7)takes the fused top-nofhybrid.search, reranks all of it and returns the firsttop_kof MMR. The tests check BM25 against hand-computed scores, that a rare term outranks a common one, that MMR drops a near-duplicate of the first pick (two chunks sharing two of three paragraphs) wherelam=1.0keeps it, the cosine fallback, that the cross encoder is lazy, that the pipeline stays inside the fused pool with a stable prefix, and that reranking raises MRR on the golden set and MMR keeps it at or above hybrid alone. -
05-metrics(Docker):metrics.relevant(chunks, q)is the set of chunk ids ofq.doc_idwhose text containsq.anchor(case-insensitive, whitespace-normalised). For a ranked list of chunk ids:recall_at_k= relevant in the topk/ relevant total;precision_at_k= relevant in the topk/k;mrr= 1 / rank of the first relevant (0 when none);ndcg_at_khas binary gains and1 / log2(rank + 1)discounts, divided by the ideal DCG overmin(k, relevant total)ranks, so one relevant chunk at rank 1 scores 1.0.metrics.evaluate(conn, embedder, reranker, golden, *, k=10, n=50, mmr_lambda=0.7) -> Metricsrunsrerank.pipelineper query (reranker=Nonemeanshybrid.searchalone), averages over the queries (MRR over the topk) and fillsby_kindwith the same means plusqueriesper kind.metrics.report(m)prints one row for all queries and one per kind.python -m retrieval_lab.metrics --corpus fixtures/corpus --k 10indexes the corpus withHashEmbedderinto its own schema (--dsn, default$RETRIEVAL_LAB_DSNor the compose database on port 5433;--schema) and prints the report;--failuresprints insteadquery_id, kind, best_rankfor every golden query whose first relevant chunk is outside the topk(best_rankis its rank in the full reranked pool, ornonewhen the pool missed it). The tests check every metric against hand-computed values, the nDCG ideal, the recall@50 of the fused pool, the reference numbers below within 0.02 and the floors, the split by kind, and both CLI outputs. -
06-answer-assembly(Docker):answer.sources(conn, embedder, reranker, query, *, top_k=5)returns the pipeline’s top chunks and their reranker scores.answer.assemble(query, chunks, *, scores=None, min_support=1, floor=SUPPORT_FLOOR) -> Draftrefuses (Draft("", (), refused=True)) when fewer thanmin_supportchunks score at leastfloor(withscores=Noneevery chunk counts); otherwise its prompt tells the model to answer only from the numbered sources, to cite every sentence as[n]and to reply exactlyI cannot answer from the sources.when none supports an answer, then gives the question and the sources as[i] <chunk_id>followed by the chunk text, andcitationsholds the chunk ids in that order.answer.check(text, citations) -> Faithfulness(cited, uncited_claims, refused)splits the text into sentences (after.,!,?followed by whitespace, or a line break, but never after the abbreviationsMr.,Mrs.,Ms.,Mme.,Dr.,St.,vs.,e.g.,i.e.,etc.; a marker that opens a sentence belongs to the one before),citedare the valid source numbers used,uncited_claimscounts the sentences without a valid marker (the refusal sentence excepted), andrefusedis whether the text contains the refusal. A number outside1..nis not a citation: its sentence counts as uncited andanswer.invalid_citations(text, citations)lists it.answer.run(query, draft, *, query_fn=None, model=None) -> Answer(text, faithfulness, cost_usd)returns the refusal without a model call for a refused draft and otherwise sends the prompt throughclaude_agent_sdk.querywithtools=[],max_turns=1,setting_sources=[],max_budget_usd=0.05,strict_mcp_config=True(your own MCP servers stay out of the call) and modelclaude-opus-5-5. The tests use the scripted SDK fake: the prompt, refusal below the floor and withmin_support=2, no model call on refusal,[1]and[3]cited, uncited sentences, abbreviations that do not split,[7]among five sources, the pinned options, that six off-corpus questions refuse on the real index and that 47 of the 48 golden queries clear the floor. -
07-retrieval-error-analysis(Docker, judgment): writeRETRIEVAL_ANALYSIS.mdin your workspace root. Runpython -m retrieval_lab.metrics --corpus fixtures/corpus --k 10 --failuresto list the golden queries the pipeline misses, then read each one: the query, the chunk that holds its anchor, the top 10, and where the anchor chunk sits in the vector list, the full-text list, the fused pool and the reranked pool. Give each miss one row of a table| query_id | failure | fix |, withfailureone ofvocabulary-gap(query and passage share no words, neither list reaches the chunk),chunk-boundary(the cut separated the anchor from the words or heading that would match),ranking(the chunk is in the fused pool but fusion ranks it low),reranker-demotion(fusion had it in the top 10 and the reranker pushed it out) oranchor-ambiguity(another chunk in the top 10 answers the question, but without the anchor), andfixone ofhybrid-weights,smaller-chunks,contextual-chunk-headers,query-rewrite,better-embedder. Then write one paragraph (25 words or more) per failure kind you used, naming the kind or sitting under a heading that names it. The test computes the misses itself from your pipeline (HashEmbedder,LexicalReranker, MMR, k = 10) and checks that every miss has a row, that both columns use the vocabulary, that at least 70% of the failure kinds agree withtests/KEY.yaml(some queries accept two readings), and that the paragraphs exist. A table with the same kind on every row, or rows for other query ids, does not pass; four disagreements out of thirteen fail. -
08-live-local-models(live, Docker and the embedding models): writeFastEmbedder.embed(fastembed’sTextEmbedding(model_name=..., cache_dir=...), imported on the first non-empty call,float32rows of 384) and useCrossEncoderRerankerfrom stage 04. The offline tests check the option shapes and the cache directory, that nothing imports fastembed before it is needed, and that both classes wrap a fake fastembed correctly. The live test indexes the corpus withFastEmbedderinto a fresh 384-dimension table, evaluates hybrid alone and the full pipeline (cross-encoder rerank of the fused top 50, then MMR), writes the aggregate toruns/live-<date>.jsonin your workspace, and asserts recall@10 ≥ 0.75 and MRR ≥ 0.60 (the course’s original target was recall@10 ≥ 0.85; see the measured numbers below). It runs for about 2 minutes on a laptop CPU (indexing 653 chunks and cross-encoding 48 × 50 pairs). -
09-live-answer(live, Docker and Claude Code):answer.runwithout a fake. Three golden queries (one per kind) go throughanswer.sources,answer.assembleandclaude-opus-5-5; the test asserts that each answer either refuses or cites a chunk that holds the anchor, that no answer cites a source number it was not given, and that the three cost at most $0.20 together, and writesruns/live-<date>.json. The offline tests replay the same path with the scripted SDK.
Measured numbers (reference, HashEmbedder + full text + LexicalReranker + MMR)
On the 48 golden queries the fused top-50 pool reaches recall@50 = 0.792, which bounds every reranker over that pool. At k = 10:
| pipeline | recall@10 | precision@10 | MRR | nDCG@10 |
|---|---|---|---|---|
| hybrid (RRF) | 0.656 | 0.069 | 0.406 | 0.468 |
| + lexical rerank | 0.750 | 0.079 | 0.545 | 0.595 |
| + MMR (λ 0.7) | 0.729 | 0.077 | 0.536 | 0.583 |
By kind with rerank and MMR: lexical recall 1.000 / MRR 0.854, factual 0.938 / 0.656, conceptual 0.250 / 0.098. Every miss but one is a conceptual query: they paraphrase without the document’s vocabulary, so neither hashing vectors nor BM25 can bridge them (stage 07 classifies them; stage 08 measures real embeddings). The course’s original targets were recall@10 ≥ 0.80, MRR ≥ 0.50 and nDCG@10 ≥ 0.55; recall@10 cannot pass 0.792 over a 50-candidate pool, so its floor is the measured value rounded down to 0.05, 0.70 (MRR and nDCG keep their original floors). Stage 05 pins the numbers above within 0.02.
The support floor of stage 06, answer.SUPPORT_FLOOR = 4.0, is a BM25 score from the lexical
reranker: off-corpus questions top out between 0 and 3.2 (their stopwords still match), the
weakest golden query tops out at 3.86 (q41, which therefore refuses) and the median golden query
at about 11. A different reranker (the cross-encoder scores logits) needs its own floor.
Measured live numbers (2026-10-23)
fixtures/live_local_models_run.json (stage 08, BAAI/bge-small-en-v1.5 + Xenova/ms-marco- MiniLM-L-6-v2 + MMR, 48 queries, k = 10):
| pipeline | recall@10 | precision@10 | MRR | nDCG@10 |
|---|---|---|---|---|
| hybrid (RRF), bge-small | 0.740 | 0.077 | 0.552 | 0.598 |
| + cross-encoder + MMR | 0.781 | 0.081 | 0.648 | 0.674 |
By kind with the cross-encoder: lexical recall 0.906 / MRR 0.896, factual 1.000 / 0.805, conceptual 0.438 / 0.243 (the hashing pipeline had 0.250 / 0.098). MRR clears its floor of 0.60; recall@10 does not reach the course’s original target of 0.85: the fused bge pool holds the anchor chunk for 0.875 of the queries at 50, and the cross-encoder lifts only part of it into the top 10, mostly failing on conceptual queries. The stage 08 floor is therefore the measured recall rounded down to 0.05, 0.75, with MRR ≥ 0.60 kept. The download was 66.5 MB (quantised bge-small) and 91 MB (MiniLM cross-encoder), 18 s; the stage ran for 132 s (index 50 s, reranked evaluation 80 s).
fixtures/live_answer_run.json (stage 09): q03 (factual), q40 (lexical) and q46 (conceptual)
each answered with citations of a relevant chunk, none refused, no invalid citation, $0.099 in
total. The recorded q40 answer counts one uncited “claim”: the splitter breaks after Sofronie.,
the period of the name inside the quoted sign.
Postgres with pgvector (stages 02 to 09)
Stages 02 to 09 need Docker. The tests start their own throwaway container
(pgvector/pgvector:pg17: Postgres 17 with pgvector 0.8) on a free port of 127.0.0.1, give
each test module its own schema, and stop the container when the run ends; you do not start
anything for them. The first run pulls the image (about 650 MB on disk). The containers carry
the label miubais=m4-tests; the next test run stops any left over for longer than an hour (a
grader timeout kills pytest before it can stop its own). To stop them by hand:
docker ps -q --filter label=miubais=m4-tests | xargs docker stopFor your own experiments, docker-compose.yml runs the same image on 127.0.0.1:5433 (5432 is
often taken by a local Postgres). From the repository root:
docker compose -f course/m4-retrieval-lab/docker-compose.yml up -dpsql postgresql://postgres:miubais@127.0.0.1:5433/postgresIn Python, connect with psycopg.connect(...), run CREATE EXTENSION IF NOT EXISTS vector once,
then pgvector.psycopg.register_vector(conn) so numpy arrays go in and out as vector values.
The password is a fixed local placeholder. Stop it with docker compose -f course/m4-retrieval-lab/docker-compose.yml down; add -v to delete its volume.
Model cache (stage 08)
FastEmbedder and the cross-encoder download their ONNX models on first use (66.5 MB and 91 MB,
about 160 MB) into cache_dir, which defaults to $FASTEMBED_CACHE_PATH when set and otherwise to
~/.cache/fastembed (fastembed’s own default is a temporary directory that the system may clear).
uv run miubais doctor reports embed-model as present once both models are in that directory.
No test before stage 08 downloads a model.
Requirements by stage
| Stage | Title | Read first | Live | Requires |
|---|---|---|---|---|
01-chunker | Chunk Markdown and plain text by headings and paragraphs | Chunking and embeddings | no | - |
02-embed-and-store | A hashing embedder and a pgvector table with full text | Chunking and embeddings | no | docker |
03-hybrid-retrieval | Vector and full-text search fused with reciprocal rank fusion | Hybrid search | no | docker |
04-reranker | A lexical reranker and MMR over the fused candidates | Reranking and MMR | no | docker |
05-metrics | Recall, precision, MRR and nDCG on the golden set | Retrieval metrics and golden sets | no | docker |
06-answer-assembly | A cited answer prompt and a faithfulness check | Citations, faithfulness and agentic retrieval | no | docker |
07-retrieval-error-analysis | Retrieval error analysis of the failed golden queries (judgment stage) | Retrieval metrics and golden sets | no | docker |
08-live-local-models | Live: local embeddings and a cross-encoder on the golden set | Chunking and embeddings | yes | docker, embed-model |
09-live-answer | Live: cited answers to three golden queries | Citations, faithfulness and agentic retrieval | yes | docker, claude-auth |
Check your work
Check your work from the clone:
uv run miubais check m4-retrieval-lab
uv run miubais hint m4-retrieval-labOpened from pnpm dev on localhost with uv run miubais serve running, this panel runs them for you.