Skip to content

reranker.md

Latest commit

 

History

History
178 lines (140 loc) · 6.96 KB

reranker.md

File metadata and controls

178 lines (140 loc) · 6.96 KB
Raw

Reference — rag-reranker

Public API for the Step 2.7 two-stage reranker. Architectural notes live in architecture/reranking.md; the design anchor is ADR-0006.

Overview

The rag-reranker package ships three pieces:

Symbol Module Purpose
RerankPipeline rag_reranker.pipeline Async orchestrator running the two-stage cascade + optional MMR.
MMRConfig rag_reranker.types Configuration for the MMR diversification pass.
mmr_rerank rag_reranker.mmr Pure MMR primitive (sync, deterministic).

Concrete reranker backends ship under rag_reranker.backends behind optional extras:

Class Extra Notes
SentenceTransformersCrossEncoder sentence-transformers Local cross-encoder. Factories: bge_reranker_v2_m3() (ADR default), minilm_l6() (fast English).
CohereReranker cohere Cohere Rerank v3 hosted API.
JinaReranker jina Jina AI Reranker hosted API via aiohttp.

Usage

Minimal pipeline (no MMR, NoopReranker)

from rag_core.spi.noop import NoopReranker
from rag_reranker import RerankPipeline

pipe = RerankPipeline(
    reranker=NoopReranker(),
    hydrator=my_hydrator,         # async (ctx, refs) -> list[Chunk]
)
chunks = await pipe.rerank(ctx, "the query", refs, top_k=10)

Cross-encoder + MMR diversification

from rag_reranker import MMRConfig, RerankPipeline
from rag_reranker.backends.sentence_transformers import bge_reranker_v2_m3

pipe = RerankPipeline(
    reranker=bge_reranker_v2_m3(),
    hydrator=hybrid_backend.hydrate,
    embedder=my_embedder,         # used by MMR when chunks lack metadata['embedding']
    mmr=MMRConfig(lambda_=0.7, top_k=10),
    stage1_top_n=50,
    stage2_top_n=20,
)
chunks = await pipe.rerank(ctx, "the query", refs, top_k=10)

Hosted-API reranker (Cohere)

from rag_reranker.backends.cohere import CohereReranker

reranker = CohereReranker(api_key=os.environ["COHERE_API_KEY"], model="rerank-v3.5")
pipe = RerankPipeline(reranker=reranker, hydrator=hybrid_backend.hydrate)

ragctl rerank — local smoke

$ ragctl rerank "what is a cross-encoder?" --top-k 3
query:    'what is a cross-encoder?'
tenant:   ragctl-local
backend:  NoopReranker
mmr:      off
top_k:    3
results:  3
  #1  score=1.0000  id=chunk-0  Cross-encoder rerankers score (query, document) pairs jointly.
  #2  score=1.0000  id=chunk-1  Reciprocal rank fusion combines multiple ranked lists.
  #3  score=1.0000  id=chunk-2  MMR balances relevance against diversity in selection.

Flags:

--top-k / -k          Final number of reranked chunks to print (default 3).
--stage1-top-n        Stage-1 (fast_rerank) truncation depth (default 50).
--stage2-top-n        Stage-2 (precise_rerank) candidate depth (default 10).
--mmr / --no-mmr      Apply MMR diversification (default off).
--mmr-lambda          MMR relevance/diversity trade-off in [0, 1] (default 0.7).
--backend             noop | sentence-transformers | cohere | jina (default noop).
--model               Model id (sentence-transformers / cohere / jina backends).
--tenant / -t         Tenant id used to build the RequestContext.

--backend cohere and --backend jina require COHERE_API_KEY / JINA_API_KEY in the environment.

Internals

RerankPipeline flow

  1. Stage 1. Reranker.fast_rerank(ctx, query, candidates, top_n=stage1_top_n) narrows the candidate set. Default implementation truncates without re-scoring.
  2. Early exit. Reranker.should_early_exit(stage1_scores) — if True, hydrate the top top_k ChunkRefs and return; stage 2 and MMR are skipped. Stage-1 scores are preserved on the way out.
  3. Hydrate. hydrator(ctx, stage1_refs) turns ChunkRef\s into Chunk\s (loading content + metadata). Any *RetrievalBackend.hydrate method satisfies this protocol.
  4. Stage 2. Reranker.precise_rerank(ctx, query, chunks, top_n=stage2_top_n) runs the cross-encoder and writes scores onto each returned chunk's Chunk.score field.
  5. MMR (optional). When MMRConfig is supplied AND there's more than one candidate, mmr_rerank(...) re-orders stage-2 output for diversity. See MMRConfig for the λ knob.

Errors

Any backend or stage failure raises RerankerError (a RetrievalError subclass). The context dict carries a stage key identifying which step failed:

stage Meaning
"fast" fast_rerank raised.
"early_exit" should_early_exit raised.
"hydrate" The hydrator callable raised.
"precise" precise_rerank raised.
"mmr" MMR was requested but embeddings could not be resolved.

mmr_rerank algorithm

Carbonell & Goldstein (1998). The selection rule is:

MMR(d) = λ · rel(d, q) − (1 − λ) · max_{s ∈ S} sim(d, s)

where S is the running set of selected documents, rel is the relevance score (passed in by the caller), and sim is cosine similarity over the supplied embeddings. The first pick uses argmax(relevance) (no penalty term) so the algorithm is stable at λ = 0.

Argument Required Notes
candidates yes Treated as a set; MMR builds its own ordering.
relevance_scores yes Same length as candidates.
embeddings yes Same length as candidates; all rows same dimension.
lambda_ no Trade-off in [0, 1]. Default 0.7. 1.0 → pure relevance; 0.0 → first pick by relevance then pure diversity thereafter.
top_k no Cap at len(candidates). Default 10.

The cosine-similarity cache is lazy — most picks only need similarity to the selected set, not to every other candidate. Returned chunks are unmodified (MMR does NOT write to Chunk.score; that's the output of precise_rerank).

Extension points

Hook What you provide Where it goes
Reranker subclass precise_rerank (required), optionally fast_rerank / should_early_exit. RerankPipeline(reranker=...).
Hydrator callable async (ctx, refs) -> list[Chunk]. RerankPipeline(hydrator=...).
Embedder Only when MMR is on and chunks lack metadata['embedding']. RerankPipeline(embedder=...).
MMRConfig λ + top_k. RerankPipeline(mmr=...).

To add a backend, see the backends folder — each existing backend is a stand-alone module under 200 lines.

See also