ripvec_core/rerank.rs
1//! Cross-encoder reranker for top-K refinement.
2//!
3//! ## Why this module exists
4//!
5//! ripvec's bi-encoder retrieval (BERT or semble) embeds query and
6//! documents into a shared vector space and ranks by cosine. That's
7//! cheap to scale, but the model can't express cross-token
8//! interactions between query and document — each side is encoded
9//! independently. On natural-language and prose corpora this caps
10//! quality.
11//!
12//! A cross-encoder concatenates the pair `[CLS] query [SEP] doc [SEP]`
13//! and runs full attention across both, producing a single relevance
14//! score. Quality is meaningfully higher but cost is O(candidates),
15//! so it's used only as a reranker on the bi-encoder's top-K.
16//!
17//! ## Architecture
18//!
19//! This module is a thin orchestrator: tokenize `(query, doc)` pairs,
20//! delegate scoring to a [`RerankBackend`](crate::backend::RerankBackend)
21//! (currently [`crate::backend::cpu::CpuRerankBackend`] — same BERT
22//! trunk as the bi-encoder, plus a `Linear(hidden -> 1)` classifier
23//! head + sigmoid).
24//!
25//! Adding GPU rerankers later is mechanical: implement
26//! `RerankBackend` for Metal/CUDA/MLX, mirror `load_reranker_cpu` in
27//! `backend/mod.rs`, route through `Reranker::from_pretrained`.
28
29use anyhow::anyhow;
30use tokenizers::Tokenizer;
31
32use crate::backend::{Encoding, RerankBackend};
33
34/// Default cross-encoder model.
35/// `cross-encoder/ms-marco-MiniLM-L-12-v2` is 33MB, ~10ms per
36/// query/doc pair on CPU, NDCG@10 = 74.5 on MS MARCO dev. Picked over
37/// the smaller L-6 (22MB, NDCG 74.3) because the 4-corpus benchmark
38/// matrix showed L-12 added meaningful target-hit lift across both
39/// prose (Gutenberg) and code (Tokio) — and the ~5ms/pair extra is
40/// invisible against the indexing budget on any non-trivial corpus.
41pub const DEFAULT_RERANK_MODEL: &str = "cross-encoder/ms-marco-MiniLM-L-12-v2";
42
43/// Default cap on candidates passed to the reranker.
44///
45/// Cost is linear in candidates; 100 is the standard top-K in the
46/// retrieve-then-rerank literature. At ~5ms/pair on MiniLM-L-6 this
47/// is ~500ms total, the upper edge of interactive.
48pub const DEFAULT_RERANK_CANDIDATES: usize = 100;
49
50/// Cross-encoder reranker orchestrator.
51///
52/// Owns a `RerankBackend` (model trunk + classifier head) and the
53/// tokenizer that produced the encodings the backend expects.
54///
55/// Construct via [`Self::from_pretrained`]. Use [`score_pairs`] to
56/// rank candidate `(query, doc)` text pairs.
57///
58/// [`score_pairs`]: Self::score_pairs
59pub struct Reranker {
60 backend: Box<dyn RerankBackend>,
61 tokenizer: Tokenizer,
62}
63
64impl Reranker {
65 /// Load a cross-encoder by `HuggingFace` repo ID.
66 ///
67 /// Routes through [`crate::backend::load_reranker_cpu`] for now;
68 /// GPU paths slot in here as feature-gated branches when added.
69 /// The tokenizer is downloaded via the same `hf-hub` cache, so
70 /// multiple sub-agent MCP processes share weights through
71 /// `~/.cache/huggingface/hub/`.
72 ///
73 /// # Errors
74 ///
75 /// Returns an error if the model can't be downloaded, lacks a
76 /// classifier head (i.e., a bi-encoder was supplied by mistake),
77 /// or fails to load.
78 pub fn from_pretrained(model_repo: &str) -> crate::Result<Self> {
79 let backend = crate::backend::load_reranker_cpu(model_repo)?;
80 let tokenizer = crate::tokenize::load_tokenizer(model_repo)?;
81 Ok(Self { backend, tokenizer })
82 }
83
84 /// Score a batch of `(query, document)` pairs.
85 ///
86 /// Returns scores in `[0, 1]` (sigmoid-activated), one per input
87 /// pair, in input order. Tokenizes with a `(query, doc)` tuple so
88 /// `token_type_ids` are 0 for the query side, 1 for the doc side —
89 /// the convention BERT cross-encoders are trained on.
90 ///
91 /// # Errors
92 ///
93 /// Propagates tokenization or forward-pass errors.
94 pub fn score_pairs(&self, pairs: &[(&str, &str)]) -> crate::Result<Vec<f32>> {
95 if pairs.is_empty() {
96 return Ok(Vec::new());
97 }
98 let max_tokens = self.backend.max_tokens();
99 let encodings: crate::Result<Vec<Encoding>> = pairs
100 .iter()
101 .map(|(q, d)| {
102 let enc = self
103 .tokenizer
104 .encode((*q, *d), true)
105 .map_err(|e| crate::Error::Other(anyhow!("rerank tokenize failed: {e}")))?;
106 let len = enc.get_ids().len().min(max_tokens);
107 Ok(Encoding {
108 input_ids: enc.get_ids()[..len].iter().map(|&x| i64::from(x)).collect(),
109 attention_mask: enc.get_attention_mask()[..len]
110 .iter()
111 .map(|&x| i64::from(x))
112 .collect(),
113 token_type_ids: enc.get_type_ids()[..len]
114 .iter()
115 .map(|&x| i64::from(x))
116 .collect(),
117 })
118 })
119 .collect();
120 let encodings = encodings?;
121 self.backend.score_batch(&encodings)
122 }
123
124 /// Max sequence length supported by the underlying model.
125 #[must_use]
126 pub fn max_tokens(&self) -> usize {
127 self.backend.max_tokens()
128 }
129}
130
131#[cfg(test)]
132mod tests {
133 use super::*;
134
135 /// `Reranker::from_pretrained` works end-to-end on the default model.
136 /// Gated `--ignored` since it downloads weights from `HuggingFace`.
137 ///
138 /// Verifies the two structural claims:
139 /// 1. The cross-encoder ranks a relevant doc higher than an
140 /// irrelevant one for the same query.
141 /// 2. Scores are in `[0, 1]` (sigmoid range).
142 #[test]
143 #[ignore = "requires network + model download (~22MB)"]
144 fn loads_and_ranks_default_cross_encoder() {
145 let rr = Reranker::from_pretrained(DEFAULT_RERANK_MODEL)
146 .expect("default cross-encoder should load");
147 let scores = rr
148 .score_pairs(&[
149 (
150 "how to make pasta",
151 "Boil water, add salt, cook pasta for 8 minutes.",
152 ),
153 (
154 "how to make pasta",
155 "The mitochondria is the powerhouse of the cell.",
156 ),
157 ])
158 .expect("scoring should succeed");
159 assert_eq!(scores.len(), 2);
160 assert!(scores.iter().all(|&s| (0.0..=1.0).contains(&s)));
161 assert!(
162 scores[0] > scores[1],
163 "relevant doc ({}) should beat irrelevant ({})",
164 scores[0],
165 scores[1]
166 );
167 }
168}