Runs in your browser — no install. (Works once the repo is public.)
# --- Colab setup (skipped when running locally) ---
import os, sys
if 'google.colab' in sys.modules and not os.path.exists('data/sample_corpus.json'):
!git clone -q https://github.com/mdhabibi/llm-search-handbook.git
%cd llm-search-handbook
!pip -q install -r requirements.txt
Chapter 7 — Re-ranking (retrieve, then re-rank)¶
First-stage retrieval is fast but approximate. A cross-encoder re-reads query + document together and re-scores the top candidates for a precise final order.
Cross-encoder downloads on first run (internet once).
Setup¶
pip install sentence-transformers
# Bootstrap: locate repo root and import shared helpers
import sys, os
d = os.getcwd()
while d != os.path.dirname(d) and not os.path.exists(os.path.join(d, 'data', 'sample_corpus.json')):
d = os.path.dirname(d)
ROOT = d; sys.path.insert(0, os.path.join(ROOT, 'src'))
import numpy as np
from corpus import load_corpus, tokenize
1. First stage: retrieve candidates (dense)¶
We pull the top-k candidates cheaply with the Chapter 5 bi-encoder.
from semantic_search import SemanticSearch
from sentence_transformers import SentenceTransformer
docs = load_corpus()
bi = SentenceTransformer('all-MiniLM-L6-v2')
engine = SemanticSearch(bi.encode).index(docs)
query = 'what should I do about a pounding headache?'
candidates = engine.search(query, k=6) # first-stage top-6
print('First-stage (bi-encoder) order:')
for i,score,doc in candidates:
print(f' {score:.3f} D{i}: {doc["title"]}')
2. Second stage: re-rank with a cross-encoder¶
The cross-encoder scores each (query, document) pair by reading them jointly.
from sentence_transformers import CrossEncoder
reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
pairs = [(query, doc['text']) for _,_,doc in candidates]
ce_scores = reranker.predict(pairs)
reranked = sorted(zip(candidates, ce_scores), key=lambda x: -x[1])
print('Second-stage (cross-encoder) order:')
for (i,_,doc), s in reranked:
print(f' {s:6.2f} D{i}: {doc["title"]}')
print('\nNote how the ordering sharpens vs. stage 1.')
3. A reusable retrieve-then-rerank function¶
The standard two-stage pattern in a few lines.
def retrieve_then_rerank(query, first_k=20, final_k=3):
cands = engine.search(query, k=first_k) # wide & cheap
scores = reranker.predict([(query, d['text']) for _,_,d in cands]) # narrow & precise
ranked = sorted(zip(cands, scores), key=lambda x: -x[1])[:final_k]
return [(i, float(s), d) for (i,_,d), s in ranked]
for i, s, d in retrieve_then_rerank('a large animal that lives in the sea'):
print(f' {s:6.2f} D{i}: {d["title"]}')
4. The cost knob: how many candidates?¶
Re-ranking is bounded by first-stage recall: the cross-encoder can only reorder what it's given. Retrieve too few and the best doc may be excluded; too many and latency grows. Typical: retrieve 50-200, re-rank to 5-10.
import time
for first_k in [3, 10, 25]:
t0=time.time(); _ = retrieve_then_rerank('a large animal that lives in the sea', first_k=first_k)
print(f'first_k={first_k:>3} -> {(time.time()-t0)*1000:6.0f} ms (more candidates = safer recall, slower)')
5. Re-ranking rescues a weak first stage¶
Re-ranking isn't only for dense retrieval. It works on any ranked list -- including plain keyword search, which has no idea what a question means.
The strategy is counter-intuitive but powerful: if your first stage is cheap, don't try to make it precise. Make it wide. Let BM25 return a large, sloppy candidate set, and let the cross-encoder find the answer inside it.
from collections import Counter, defaultdict
import math
# A compact BM25 over the corpus (same scoring as Chapter 2).
tok_docs = [tokenize(d['text']) for d in docs]
avgdl = sum(len(t) for t in tok_docs) / len(tok_docs)
df = defaultdict(int)
for t in tok_docs:
for term in set(t):
df[term] += 1
def idf(term):
n = len(tok_docs)
return math.log(1 + (n - df[term] + 0.5) / (df[term] + 0.5))
def bm25_search(query, k=5, k1=1.5, b=0.75):
q = tokenize(query)
scored = []
for i, t in enumerate(tok_docs):
f = Counter(t)
s = sum(idf(term) * (f[term] * (k1 + 1)) /
(f[term] + k1 * (1 - b + b * len(t) / avgdl))
for term in q if term in f)
scored.append((i, s))
return sorted(scored, key=lambda x: -x[1])[:k]
kw_query = 'what should I do about a pounding headache?'
print('BM25 top-3 (keyword only):')
for i, s in bm25_search(kw_query, k=3):
print(f' {s:5.2f} D{i}: {docs[i]["title"]}')
Keyword search matches words, so it rewards passages that happen to repeat query terms — it cannot tell which passage actually answers the question.
Now widen the net and let the cross-encoder do the judging.
# Stage 1: cast a WIDE net with cheap keyword search (no relevance judgement at all)
wide = bm25_search(kw_query, k=len(docs)) # in production: 100-500
wide_ids = [i for i, _ in wide]
# Stage 2: the cross-encoder reads each (query, passage) pair and scores relevance
scores = reranker.predict([(kw_query, docs[i]['text']) for i in wide_ids])
rescued = sorted(zip(wide_ids, scores), key=lambda x: -x[1])[:3]
print('After re-ranking the wide keyword list:')
for i, s in rescued:
was = wide_ids.index(i) + 1
print(f' {s:6.2f} D{i}: {docs[i]["title"]:<32} (was rank {was} in BM25)')
Watch the was rank N column: the re-ranker pulls the right passage up from deep in a list
that keyword search had ordered badly.
The trade-off. A wider candidate set raises your ceiling (the answer is more likely to be somewhere in the list) but costs one cross-encoder pass per candidate. That's the same recall-vs-latency knob from section 4 -- now applied to a first stage that is nearly free.
This is also the honest argument for hybrid search (Chapter 8): rather than fixing a weak first stage with a big k, fuse BM25 and dense retrieval, then re-rank a smaller, better candidate set.
Takeaway & exercises¶
Retrieve wide and cheap (bi-encoder), then re-rank narrow and precise (cross-encoder). This two-stage pattern is the backbone of production search and RAG.
Exercises 0. Ask something where similarity misleads -- e.g. a query whose nearest passage is on-topic but answers a different question. Does the re-ranker demote it?
- Find a query where stage-1 rank #1 is wrong but re-ranking fixes it.
- Measure how final quality changes as
first_kgrows from 3 to 15. - Add a 'prefer shorter passages' tie-breaker to the re-rank sort. Where in the pipeline does it belong?