Implement a simple re-ranking model for retrieved results.

What you need to know

Two-stage retrieval gets the best of both: cheap recall first, expensive precision second. The retriever's job is to make sure the right answer is somewhere in the top 50; the re-ranker's job is to put it at position 1.

The re-rank function should not care which scorer it gets. It takes a score_pairs function, so in the interview you can plug in a trained cross-encoder, an LLM, or a toy scorer for testing.

1from collections.abc import Callable, Sequence23ScoreFn = Callable[[list[tuple[str, str]]], Sequence[float]]45def rerank(query: str, candidates: list[str], score_pairs: ScoreFn,6           top_k: int = 5, batch_size: int = 32) -> list[tuple[str, float]]:7    """Score (query, candidate) pairs in batches and return the top_k, best first."""8    candidates = [c for c in candidates if c and c.strip()]9    if not candidates:10        return []11    pairs = [(query, c) for c in candidates]12    scores: list[float] = []13    for i in range(0, len(pairs), batch_size):14        scores.extend(float(s) for s in score_pairs(pairs[i:i + batch_size]))15    ranked = sorted(zip(candidates, scores), key=lambda p: p[1], reverse=True)16    return ranked[:top_k]

With a real model (the sentence-transformers library):

1from sentence_transformers import CrossEncoder23_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")45def cross_encoder_scores(pairs: list[tuple[str, str]]) -> list[float]:6    return _model.predict(pairs).tolist()

The tricky parts:

• Batching sends 32 pairs per model call. One pair per call wastes GPU time; all pairs in one call can run out of memory.
• sorted is stable, so candidates with equal scores keep the retriever's order — a sensible tie-break for free.
• Empty and whitespace strings are dropped before scoring; a model scores them anyway and they can land in the top 5.

Complexity: m candidates means m transformer passes, each roughly O(L²) in the pair's token length L because of attention. The sort is O(m log m), which is nothing next to the model. Space is O(m). Latency grows linearly with m, so keep m between about 25 and 100.

A real-life example

A toy scorer that counts shared words stands in for the cross-encoder so the numbers are checkable:

1def overlap_scores(pairs):2    return [float(len(set(q.lower().split()) & set(d.lower().split()))) for q, d in pairs]34candidates = ["API rate limits for free plans",          # retriever rank 15              "Rotate an API key from Settings, then Keys",6              "Billing keys and invoices",7              "How to rotate a password"]8print(rerank("how do i rotate an api key", candidates, overlap_scores, top_k=2))9# [('Rotate an API key from Settings, then Keys', 4.0), ('How to rotate a password', 2.0)]

The retriever had put "API rate limits" first because it is close in topic; the re-ranker, looking at both texts together, moves the real answer to the top. Run with the real cross_encoder_scores above, the MiniLM cross-encoder agrees: the Settings page scores about 7.8 and the password page about -3.7 (these are raw relevance logits, not probabilities).

Developer-docs search, e-commerce search and every serious RAG stack use this retrieve-50, re-rank-to-5 shape.

Follow-up questions to expect

• "No model available in the interview — what then?" — Use an LLM as the scorer: send the query and all candidates with ids in one call and ask for the ids in relevance order. One call is much cheaper than one call per document; validate that every returned id exists.
• "How do you know the re-ranker helps?" — Measure nDCG@5 or recall@5 on a labelled set before and after. If it does not move, it is only adding latency.
• "Documents longer than the model's 512-token limit?" — The model truncates silently, and the relevant passage may be cut off. Re-rank chunks, not whole documents.