Runs in your browser — no install. (Works once the repo is public.)

In [ ]:

# --- 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

# --- 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

In [ ]:

# 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

# 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.

In [ ]:

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"]}')

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.

In [ ]:

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.')

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.

In [ ]:

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"]}')

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.

In [ ]:

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)')

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)')

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

1. Find a query where stage-1 rank #1 is wrong but re-ranking fixes it.
2. Measure how final quality changes as first_k grows from 3 to 15.
3. Add a 'prefer shorter passages' tie-breaker to the re-rank sort. Where in the pipeline does it belong?