Retrieval quality & re-ranking

1 · Why the first pass isn't enough

First-pass retrieval — vector search or BM25 — is built for speed over a huge corpus. It embeds each document once and compares vectors independently, so it scales to millions of chunks. The price is precision: it finds things that are topically near the question, but it can't tell that chunk #5 actually answers the question while chunk #1 merely mentions the same words. That is a precision@k problem — the right chunk is somewhere in the top 20, just not at the top.

2 · Re-ranking with a cross-encoder

A bi-encoder (what first-pass vector search uses) embeds the query and each chunk separately, then compares the two vectors. It's fast because chunk vectors are precomputed — but the model never sees the query and the chunk together. A cross-encoder does the opposite: it feeds the pair (query, chunk) through the model jointly and outputs a single relevance score. That joint view is far more accurate, but it can't be precomputed and is far slower — so you only run it on the top-N candidates the bi-encoder already found.

Below is the whole idea, runnable. First-pass retrieval is a simple shared-term count (standing in for the bi-encoder) — which a keyword-stuffed chunk can game. Then a richer re-ranker rescores the top-N with signals the first pass ignored — term coverage plus exact phrase/bigram matches — and re-orders them. Watch the truly-relevant chunk climb from behind a noisy one to the top.

rerank.pyimport re
from collections import Counter

def toks(s):
    return re.findall(r'[a-z]+', s.lower())

# ---- FIRST PASS: cheap shared-term count (stands in for a bi-encoder) ----
def first_pass_score(query, chunk):
    q, c = Counter(toks(query)), Counter(toks(chunk))
    return sum(min(q[w], c[w]) for w in q)   # raw shared-term count

# ---- RE-RANKER: richer joint score (stands in for a cross-encoder) ----
# A real cross-encoder replaces this function; the pipeline shape is identical.
def rerank_score(query, chunk):
    cl = chunk.lower()
    qset = set(toks(query))
    coverage = sum(1 for w in qset if w in set(toks(chunk))) / (len(qset) or 1)
    bigram = 1.0 if 'consumer group' in cl else 0.0   # meaningful phrase, not just words
    phrase = 1.0 if 'restart checkout' in cl else 0.0  # answers the actual ask
    return round(2.0*coverage + bigram + phrase, 3)

chunks = [
    'Restart the queue, restart the group, restart the checkout queue consumer.',  # keyword-stuffed noise
    'To restart checkout, drain the queue then scale the consumer group up.',       # the real answer
    'Checkout reads orders; the queue and consumer group appear in Grafana.',
    'Never restart the database to fix a queue backlog.',
]
query = 'restart checkout consumer group queue'

# stage 1: first-pass top-N (cheap, favors recall)
N = 4
first = sorted(chunks, key=lambda c: first_pass_score(query, c), reverse=True)[:N]
print('FIRST PASS (top-4):')
for i, c in enumerate(first, 1):
    print(f'  {i}. [{first_pass_score(query, c)}] {c}')

# stage 2: re-rank ONLY those N (richer, favors precision)
reranked = sorted(first, key=lambda c: rerank_score(query, c), reverse=True)
print('\nAFTER RE-RANK:')
for i, c in enumerate(reranked, 1):
    print(f'  {i}. [{rerank_score(query, c)}] {c}')

FIRST PASS (top-4):
  1. [5] Restart the queue, restart the group, restart the checkout queue consumer.
  2. [5] To restart checkout, drain the queue then scale the consumer group up.
  3. [4] Checkout reads orders; the queue and consumer group appear in Grafana.
  4. [2] Never restart the database to fix a queue backlog.

AFTER RE-RANK:
  1. [4.0] To restart checkout, drain the queue then scale the consumer group up.
  2. [2.6] Checkout reads orders; the queue and consumer group appear in Grafana.
  3. [2.0] Restart the queue, restart the group, restart the checkout queue consumer.
  4. [0.8] Never restart the database to fix a queue backlog.

First-pass tied a keyword-stuffed noise chunk with the real answer (both score 5), and a plain sort left the noise at #1. The re-ranker — rewarding coverage plus the meaningful consumer group and restart checkout phrases rather than raw repeats — pushed the actionable answer to #1 and demoted the noise to #3. That is precision@k improving in one sort — and we only paid the richer cost on 4 candidates, not the whole corpus.

3 · Query transformation

The other way to improve retrieval is to fix the query before you search. A user's question is often short, ambiguous, or phrased unlike the documents. Three published techniques rewrite it:

Here is a runnable multi-query simulation: several paraphrases each run first-pass retrieval, and we union + dedup the hits into one candidate set. In production the paraphrases come from an LLM; here we hand-write them so it runs offline.

multi_query.pyimport re
from collections import Counter

def toks(s):
    return re.findall(r'[a-z]+', s.lower())

def score(query, chunk):
    q, c = Counter(toks(query)), Counter(toks(chunk))
    return sum(min(q[w], c[w]) for w in q)

def retrieve(query, chunks, k=1):
    ranked = sorted(range(len(chunks)), key=lambda i: score(query, chunks[i]), reverse=True)
    return [i for i in ranked if score(query, chunks[i]) > 0][:k]

chunks = [
    'Restart checkout by draining the order queue.',
    'The vehicle will not crank when the battery is flat.',
    'Reboot the stuck payment worker to clear the consumer group.',
    'Invoices are stored in Postgres by the billing service.',
]

# One LLM would generate these paraphrases; we hand-write them for an offline run.
paraphrases = [
    'how to restart checkout queue',
    'reboot the payment worker consumer group',
    'clear a stuck checkout',
]

seen, union = set(), []          # dedup by chunk index, keep first-seen order
for pq in paraphrases:
    hits = retrieve(pq, chunks, k=1)
    print(f'{pq!r:45} -> chunks {hits}')
    for i in hits:
        if i not in seen:
            seen.add(i); union.append(i)

print('\nUNION (deduped):', union)
for i in union: print('  -', chunks[i])

'how to restart checkout queue'               -> chunks [0]
'reboot the payment worker consumer group'    -> chunks [2]
'clear a stuck checkout'                      -> chunks [2]

UNION (deduped): [0, 2]
  - Restart checkout by draining the order queue.
  - Reboot the stuck payment worker to clear the consumer group.

Each paraphrase, taking the single best hit, saw the question from a different angle: the first surfaced the “drain the queue” chunk, the other two both surfaced the “reboot the worker” chunk. The union gathered both good chunks that no single query returned together, and dedup kept the repeated one just once. That is multi-query in miniature: several angles on the question, merged. (RAG 3.2 covered fusing rankings with RRF; here we fuse sets from paraphrased queries — a complementary move.)

4 · Context packing — order, dedup, budget

You now have good candidates. The last decision is how to lay them into the prompt, and it is not “paste all of them.” Three forces shape the packing:

Here is a runnable packer that takes ranked chunks, drops near-duplicates, orders them so the best sit at the edges (fighting lost-in-the-middle), and truncates to a token budget. We approximate tokens as words to keep it offline; a real tokenizer swaps in without changing the shape.

pack.pyimport re

def approx_tokens(s):
    return len(s.split())          # a real tokenizer replaces this

def near_dup(a, b, thresh=0.8):
    wa, wb = set(re.findall(r'[a-z]+', a.lower())), set(re.findall(r'[a-z]+', b.lower()))
    if not wa or not wb: return False
    jaccard = len(wa & wb) / len(wa | wb)
    return jaccard >= thresh

def pack(ranked, budget=40):
    # ranked: chunks already sorted best -> worst (e.g. from the re-ranker)
    kept = []
    for c in ranked:                                   # 1) dedup near-duplicates
        if not any(near_dup(c, k) for k in kept):
            kept.append(c)
    trimmed, used = [], 0                              # 2) truncate to token budget
    for c in kept:
        t = approx_tokens(c)
        if used + t > budget: break
        trimmed.append(c); used += t
    # 3) order: best first, 2nd-best LAST, rest in the middle (fight lost-in-middle)
    ordered = trimmed[:]
    if len(ordered) >= 3:
        ordered = [trimmed[0]] + trimmed[2:] + [trimmed[1]]
    return ordered, used

ranked = [
    'Drain the queue then scale the consumer group to restart checkout.',
    'Drain the queue then scale the consumer group to restart the checkout.',  # near-dup of #0
    'Never restart the database to fix a queue backlog.',
    'Checkout orders flow through the queue into the consumer group.',
    'Grafana dashboards track queue depth and consumer lag over time.',
]
ordered, used = pack(ranked, budget=40)
print(f'packed {len(ordered)} chunks, ~{used} tokens:')
for i, c in enumerate(ordered, 1): print(f'  [{i}] {c}')

packed 4 chunks, ~40 tokens:
  [1] Drain the queue then scale the consumer group to restart checkout.
  [2] Checkout orders flow through the queue into the consumer group.
  [3] Grafana dashboards track queue depth and consumer lag over time.
  [4] Never restart the database to fix a queue backlog.

The near-duplicate second chunk (identical words) was dropped, the packer filled right up to the 40-token budget, and the second-strongest chunk was moved to the end so the two most important chunks sit at the edges where the model reads best. Same facts, no redundancy, better placement — the middle holds the weaker context that matters least if the model skims it.

5 · When re-ranking is worth the latency

Re-ranking and query transformation both cost time — an extra model pass, sometimes an extra LLM call. That's a real budget you spend on every query. Spend it where it pays off: