Lab 01 — Rerank: The Cross-Encoder Second Stage

Phase 31 · Lab 01 · Phase README · Warmup

The problem

First-stage retrieval — dense embeddings, BM25, or the hybrid of both you built in Phase 05 — is a bi-encoder architecture: the query and every document are embedded separately, so document vectors can be precomputed and indexed, and query time is one embedding plus a nearest-neighbor lookup. Fast, scalable — and structurally blind. Because the query never "sees" the document during scoring, a bi-encoder cannot tell "reset the vpn password on the corporate gateway" (the exact procedure) from a policy page that merely contains the words reset, vpn, password, corporate, and gateway scattered across three sentences. Word order, phrasing, and term proximity are invisible to a dot product over independently-built vectors.

Cohere's Rerank API is the productized fix: a cross-encoder that scores each (query, document) pair jointly, run only on the ~top-N candidates the cheap first stage already shortlisted. Retrieve wide and cheap; rerank narrow and precise. It is Cohere's most-loved API precisely because it drops into any existing search stack — Elasticsearch, OpenSearch, pgvector, a hand-rolled hybrid — as one extra call that visibly improves the top of the ranking.

What you build

File map

Run it

pytest test_lab.py -v                       # your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # the reference (green)
python solution.py                          # worked example

Success criteria

• relevance_score rewards exact phrasing over the same words scrambled, and tight term clusters over the same terms scattered — the two signals a bi-encoder cannot see.
• rerank returns the same candidate set in a different order than the first-stage ranking, scores descending, ties broken by original index.
• The planted dHard document (first-stage rank #4, exact-phrase answer) is promoted to reranked rank #1.
• top_n truncates the reranked list as a prefix; every result's index maps back to the input list; boosting first-stage scores changes nothing about rerank output.
• With context_window set, a relevant passage past the first chunk scores 0.0 at max_chunks_per_doc=1 and is recovered at a higher chunk budget; the doc takes its best chunk's score, not an average.
• Empty documents list → []; empty query → all-zero scores in original order; top_n <= 0 and bad chunking params raise ValidationError.
• All 27 tests pass under both lab and solution.

How this maps to the real stack

• Cohere's rerank endpoint takes a query, a list of documents (strings or objects with a text field), and top_n, and returns results carrying index (position in your input list), relevance_score (a calibrated [0, 1] relevance), and optionally the document — the exact response shape RerankResult mirrors. The index field is load-bearing in production: your candidates carry metadata (doc ids, ACLs, first-stage scores) the reranker never sees, and index is how you rejoin them.
• The model behind it (rerank-v3.5 and the earlier rerank-english / rerank-multilingual v3 models) is a trained cross-encoder: query and document are concatenated into one input and scored by a transformer that attends across them — real query×document token interaction, of which our coverage/phrase/proximity blend is the deterministic, inspectable stand-in. Open-weight equivalents (bge-reranker, mixedbread, MiniLM cross-encoders) have the same shape.
• Long documents: Cohere Rerank enforces a per-document token limit and (in earlier API versions, via max_chunks_per_doc) splits long documents into chunks, scoring the document by its best chunk — exactly the context_window/max_chunks_per_doc behaviour here, including the failure mode: at one chunk, an answer buried past the window is invisible.
• Where it sits in a stack: after hybrid first-stage retrieval (Phase 05's dense+BM25+RRF) and before generation. The latency/quality tradeoff is the design decision: a cross-encoder is O(candidates) model inferences per query, so you never run it over the corpus — you run it over the shortlist. Typical shape: retrieve 50–100, rerank, keep the top 3–10 for the context window.

Limits of the miniature. A real cross-encoder is a trained transformer that understands paraphrase and negation — our lexical interaction score rewards phrasing but cannot know that "credentials rotation" answers "password reset." Real relevance scores are model-calibrated, not a fixed weighted blend. And real rerankers score chunks in parallel on GPUs; ours is a Python loop, which is fine because the decision structure (pair-scoring, best-chunk, top_n, index preservation) is the lesson.

Extensions (your own machine)

• Swap in a real cross-encoder via sentence-transformers (cross-encoder/ms-marco-MiniLM-L-6-v2) behind the same rerank signature and compare orderings against the lexical stand-in.
• Add a return_documents=False mode that omits document text from results (Cohere's bandwidth optimization for large candidate sets).
• Measure the two-stage win properly: recall@k and MRR (Phase 05's metrics) for first-stage-only vs retrieve→rerank over a labeled query set.
• Add a latency budget model: given per-pair scoring cost, find the largest candidate count that keeps rerank under a p95 target.

Interview / resume signal

"Built a miniature of Cohere's Rerank stage: a deterministic cross-encoder-style scorer (term coverage + phrase match + proximity window) over first-stage candidates, with Cohere's exact response contract (original-index preservation, top_n truncation, calibrated-descending scores) and max_chunks_per_doc-style long-document chunking — and a planted corpus proving the two-stage win: a document buried at first-stage rank #4 promoted to #1 by the reranker."