post_cortex_embeddings/rerank/
onnx.rs

1// Copyright (c) 2026 Julius ML
2// Licensed under the MIT License. See LICENSE at the workspace root.
3
4//! BERT cross-encoder on ONNX Runtime with int8 weights: the query and a
5//! document go through the encoder together as one `[CLS] query [SEP]
6//! document [SEP]` sequence and the model outputs one relevance logit. The
7//! `cross-encoder/ms-marco-*` repositories ship dynamically quantised
8//! exports under `onnx/`, one per instruction set.
9
10use std::path::Path;
11use std::sync::{Arc, OnceLock};
12
13use anyhow::{Result, anyhow};
14use async_trait::async_trait;
15use hf_hub::api::tokio::Api;
16use ort::environment::Environment;
17use ort::session::Session;
18use ort::session::builder::{AutoDevicePolicy, GraphOptimizationLevel};
19use ort::value::Tensor;
20use parking_lot::Mutex;
21use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer, TruncationParams, TruncationStrategy};
22use tracing::info;
23
24use super::Reranker;
25
26/// Pairs per run. ONNX Runtime parallelises inside one run, so one batch
27/// per call would do; batches bound the padding to the longest pair in each.
28const BATCH_SIZE: usize = 64;
29
30struct Model {
31    session: Mutex<Session>,
32    tokenizer: Tokenizer,
33    /// Whether the graph takes `token_type_ids` (BERT does, some exports not).
34    token_types: bool,
35}
36
37/// Local int8 cross-encoder reranker on ONNX Runtime.
38pub struct OnnxCrossEncoder {
39    model: Arc<Model>,
40}
41
42/// ONNX Runtime allows one environment per process; every session shares it.
43fn environment() -> Result<&'static Environment> {
44    static ENV: OnceLock<Result<Environment, String>> = OnceLock::new();
45    ENV.get_or_init(|| init().map_err(|e| format!("{e:#}")))
46        .as_ref()
47        .map_err(|e| anyhow!("ONNX Runtime init failed: {e}"))
48}
49
50/// Load `libonnxruntime.so` from `ORT_DYLIB_PATH` (default: the library
51/// search path). Done explicitly so a missing library is an error, not the
52/// panic `ort` raises when it loads the library implicitly.
53#[cfg(target_os = "linux")]
54fn init() -> Result<Environment> {
55    let path = std::env::var("ORT_DYLIB_PATH")
56        .ok()
57        .filter(|p| !p.is_empty())
58        .unwrap_or_else(|| "libonnxruntime.so".to_string());
59    let builder = ort::init_from(&path).map_err(|e| anyhow!("loading {path}: {e}"))?;
60    builder.build().map_err(|e| anyhow!("{e}"))
61}
62
63#[cfg(not(target_os = "linux"))]
64fn init() -> Result<Environment> {
65    ort::init().build().map_err(|e| anyhow!("{e}"))
66}
67
68/// The int8 export in the model repository that fits this CPU.
69fn quantized_file() -> &'static str {
70    #[cfg(target_arch = "aarch64")]
71    {
72        "onnx/model_qint8_arm64.onnx"
73    }
74    #[cfg(target_arch = "x86_64")]
75    {
76        if std::arch::is_x86_feature_detected!("avx512vnni") {
77            "onnx/model_qint8_avx512_vnni.onnx"
78        } else if std::arch::is_x86_feature_detected!("avx512f") {
79            "onnx/model_qint8_avx512.onnx"
80        } else {
81            "onnx/model_quint8_avx2.onnx"
82        }
83    }
84    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
85    {
86        "onnx/model.onnx"
87    }
88}
89
90impl OnnxCrossEncoder {
91    /// Download (if needed) the int8 export for this CPU and the tokenizer
92    /// from the HuggingFace Hub. Each pair is truncated to `max_length`
93    /// tokens, document side first.
94    pub async fn load(model_id: &str, max_length: usize) -> Result<Self> {
95        let file = quantized_file();
96        info!("Loading ONNX cross-encoder {model_id} ({file}, max_length {max_length})");
97        let repo = Api::new()?.model(model_id.to_string());
98        let weights = repo.get(file).await?;
99        let tokenizer = repo.get("tokenizer.json").await?;
100        Self::from_files(&weights, &tokenizer, max_length)
101    }
102
103    /// Load from a local `.onnx` graph and `tokenizer.json`.
104    pub fn from_files(weights: &Path, tokenizer: &Path, max_length: usize) -> Result<Self> {
105        let mut tokenizer = Tokenizer::from_file(tokenizer).map_err(|e| anyhow!(e))?;
106        tokenizer
107            .with_truncation(Some(TruncationParams {
108                max_length,
109                strategy: TruncationStrategy::OnlySecond,
110                ..Default::default()
111            }))
112            .map_err(|e| anyhow!(e))?;
113        tokenizer.with_padding(Some(PaddingParams {
114            strategy: PaddingStrategy::BatchLongest,
115            ..Default::default()
116        }));
117
118        // `ort` defaults sessions to the MaxEfficiency device policy, which
119        // ran this model ~5x slower than ONNX Runtime's own defaults (2.1 s
120        // vs 0.43 s for 300 pairs on an M1 Pro); ask for the CPU at full speed.
121        let session = Session::builder(environment()?)
122            .map_err(|e| anyhow!("{e}"))?
123            .with_auto_device(AutoDevicePolicy::PreferCPU)
124            .map_err(|e| anyhow!("{e}"))?
125            .with_optimization_level(GraphOptimizationLevel::Level3)
126            .map_err(|e| anyhow!("{e}"))?
127            .with_intra_threads(num_cpus::get_physical())
128            .map_err(|e| anyhow!("{e}"))?
129            .commit_from_file(weights)
130            .map_err(|e| anyhow!("{e}"))?;
131        let token_types = session
132            .inputs()
133            .iter()
134            .any(|i| i.name() == "token_type_ids");
135
136        Ok(Self {
137            model: Arc::new(Model {
138                session: Mutex::new(session),
139                tokenizer,
140                token_types,
141            }),
142        })
143    }
144}
145
146impl Model {
147    fn score_batch(&self, query: &str, documents: &[String]) -> Result<Vec<f32>> {
148        let pairs: Vec<(String, String)> = documents
149            .iter()
150            .map(|d| (query.to_string(), d.clone()))
151            .collect();
152        let encodings = self
153            .tokenizer
154            .encode_batch(pairs, true)
155            .map_err(|e| anyhow!(e))?;
156        let rows = encodings.len();
157        let cols = encodings.first().map_or(0, |e| e.len());
158
159        let mut ids = Vec::with_capacity(rows * cols);
160        let mut types = Vec::with_capacity(rows * cols);
161        let mut mask = Vec::with_capacity(rows * cols);
162        for e in &encodings {
163            ids.extend(e.get_ids().iter().map(|&x| i64::from(x)));
164            types.extend(e.get_type_ids().iter().map(|&x| i64::from(x)));
165            mask.extend(e.get_attention_mask().iter().map(|&x| i64::from(x)));
166        }
167        let shape = [rows, cols];
168        let mut inputs = vec![
169            ("input_ids", Tensor::from_array((shape, ids))?.into_dyn()),
170            (
171                "attention_mask",
172                Tensor::from_array((shape, mask))?.into_dyn(),
173            ),
174        ];
175        if self.token_types {
176            inputs.push((
177                "token_type_ids",
178                Tensor::from_array((shape, types))?.into_dyn(),
179            ));
180        }
181
182        let mut session = self.session.lock();
183        let outputs = session.run(inputs)?;
184        let (_, logits) = outputs[0].try_extract_tensor::<f32>()?;
185        // Relevance in [0, 1]: the sigmoid of the logit.
186        // Logits are [rows, 1].
187        Ok(logits.iter().map(|&l| 1.0 / (1.0 + (-l).exp())).collect())
188    }
189
190    fn score(&self, query: &str, documents: &[String]) -> Result<Vec<f32>> {
191        let mut scores = Vec::with_capacity(documents.len());
192        for chunk in documents.chunks(BATCH_SIZE) {
193            scores.extend(self.score_batch(query, chunk)?);
194        }
195        Ok(scores)
196    }
197}
198
199#[async_trait]
200impl Reranker for OnnxCrossEncoder {
201    async fn score(&self, query: &str, documents: &[String]) -> Result<Vec<f32>> {
202        if documents.is_empty() {
203            return Ok(Vec::new());
204        }
205        let model = Arc::clone(&self.model);
206        let query = query.to_string();
207        let documents = documents.to_vec();
208        tokio::task::spawn_blocking(move || model.score(&query, &documents)).await?
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use std::time::Instant;
215
216    use super::*;
217
218    const MODEL: &str = "cross-encoder/ms-marco-MiniLM-L4-v2";
219
220    #[tokio::test]
221    #[ignore = "requires network + HF Hub download (~25MB)"]
222    async fn ranks_the_answer_first() {
223        let ce = OnnxCrossEncoder::load(MODEL, 256).await.unwrap();
224        let docs = vec![
225            "Melanie: I went camping with the kids last weekend.".to_string(),
226            "Caroline: I researched adoption agencies yesterday.".to_string(),
227            "Melanie: The weather has been great lately.".to_string(),
228        ];
229        let scores = ce
230            .score("What did Caroline research?", &docs)
231            .await
232            .unwrap();
233        assert_eq!(scores.len(), 3);
234        assert!(scores[1] > scores[0] && scores[1] > scores[2], "{scores:?}");
235    }
236
237    /// Latency for a 300-candidate pool of LoCoMo-length turns. Run with
238    /// `cargo test --release --features rerank -- --ignored --nocapture onnx_pool_latency`.
239    #[tokio::test]
240    #[ignore = "benchmark; requires network + HF Hub download"]
241    async fn onnx_pool_latency() {
242        let ce = OnnxCrossEncoder::load(MODEL, 256).await.unwrap();
243        let docs: Vec<String> = (0..300)
244            .map(|i| {
245                format!(
246                    "8 May 2023 · Melanie Melanie: turn {i}, we talked about painting a lake \
247                     sunrise and the pottery class I signed up for with the kids last week."
248                )
249            })
250            .collect();
251        let query = "When did Melanie paint a sunrise?";
252        ce.score(query, &docs).await.unwrap();
253        let mut times = Vec::new();
254        for _ in 0..5 {
255            let t = Instant::now();
256            ce.score(query, &docs).await.unwrap();
257            times.push(t.elapsed().as_millis());
258        }
259        times.sort_unstable();
260        println!("onnx int8, 300 pairs: p50 {} ms, all {times:?}", times[2]);
261    }
262}
post_cortex_embeddings/rerank/onnx.rs

post_cortex_embeddings/rerank/
onnx.rs
