feat: add nanobeir bench

Cargo.toml

@@ -63,3 +63,7 @@ harness = false
 [[bench]]
 name = "quantize"
 harness = false
+
+[[bench]]
+name = "nanobeir"
+harness = false

benches/README.md

@@ -1,5 +1,7 @@
 # Benchmarks
 
+- `nanobeir.rs`: retrieval quality (NDCG@10) of one model on NanoBEIR.
+  Normally run through `scripts/nanobeir.py`, not directly.
 - `encode.rs`: latency percentiles of `StaticModel::encode` for one revision.
 - `quantize.rs`: throughput and accuracy of `StaticModel::quantize_int8`
   against the f32 model, in one process (see below).
@@ -70,3 +72,28 @@ variant follows every other equally often. It reports throughput, the
 speedup over f32 with a 95% bootstrap CI, and cosine similarity to the f32
 output. Set `RAYON_NUM_THREADS=1` for single-thread numbers.
 
+
+## Retrieval quality (NanoBEIR)
+
+Run this before every release:
+
+```sh
+uv run scripts/nanobeir.py
+```
+
+The script downloads the 13 NanoBEIR datasets (`zeta-alpha-ai/Nano*`) into
+`target/nanobeir/data/`, builds `nanobeir.rs`, and scores every model in
+`benches/nanobeir_baseline.json` (NDCG@10, cosine similarity, binary
+relevance, scored like `pytrec_eval`). It fails if a model's mean score
+drops by more than `--tolerance` (default 0.002) or any single dataset
+drops by more than `--dataset-tolerance` (default 0.01). Scoring all
+baseline models takes about fifteen seconds once the data and models are
+cached.
+
+A model suffixed with `:int8` is scored after `StaticModel::quantize_int8`;
+one suffixed with `:model2vec-int8` is first quantized to int8 by Python
+model2vec (pinned in the script) and saved under `target/nanobeir/models/`.
+To score other models, pass `--models`. When a change is meant to alter
+scores, review the deltas and record them with `--update-baseline`
+(optionally restricted with `--models`). The baseline scores match Python
+model2vec scored with `pytrec_eval` to within 2e-4.

benches/nanobeir.rs

@@ -0,0 +1,135 @@
+// NanoBEIR retrieval quality (NDCG@10) of one model, driven by
+// `scripts/nanobeir.py`, which downloads the data and gates releases on the
+// scores. Uses only the public API so the same file builds against any
+// revision.
+//
+//   cargo bench --bench nanobeir -- data=target/nanobeir/data model=minishlab/potion-base-8M \
+//       [quantize=int8]
+//
+// `data` holds one directory per dataset with `corpus.jsonl` and
+// `queries.jsonl` (`{"_id", "text"}`) and `qrels.jsonl`
+// (`{"query-id", "corpus-id"}`, binary relevance). Documents are ranked by
+// cosine similarity; a document with the same id as the query is never
+// retrieved, as in BEIR, and ties go to the larger document id, as in
+// trec_eval. Prints one `SCORE <dataset> <ndcg@10>` line per
+// dataset.
+
+use std::collections::{HashMap, HashSet};
+use std::env;
+use std::fs;
+use std::path::Path;
+
+use model2vec_rs::model::StaticModel;
+use serde_json::Value;
+
+const K: usize = 10;
+
+fn read_jsonl(path: &Path) -> Vec<Value> {
+    fs::read_to_string(path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()))
+        .lines()
+        .filter(|l| !l.trim().is_empty())
+        .map(|l| serde_json::from_str(l).unwrap_or_else(|e| panic!("bad line in {}: {e}", path.display())))
+        .collect()
+}
+
+fn field(v: &Value, key: &str) -> String {
+    v[key]
+        .as_str()
+        .unwrap_or_else(|| panic!("missing string field {key:?} in {v}"))
+        .to_string()
+}
+
+fn read_texts(path: &Path) -> (Vec<String>, Vec<String>) {
+    read_jsonl(path)
+        .iter()
+        .map(|v| (field(v, "_id"), field(v, "text")))
+        .unzip()
+}
+
+fn normalized(mut v: Vec<f32>) -> Vec<f32> {
+    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
+    if norm > 0.0 {
+        v.iter_mut().for_each(|x| *x /= norm);
+    }
+    v
+}
+
+fn ndcg_at_k(ranked: &[usize], doc_ids: &[String], relevant: &HashSet<String>) -> f64 {
+    let dcg: f64 = ranked
+        .iter()
+        .enumerate()
+        .filter(|(_, d)| relevant.contains(&doc_ids[**d]))
+        .map(|(rank, _)| 1.0 / (rank as f64 + 2.0).log2())
+        .sum();
+    let idcg: f64 = (0..relevant.len().min(K))
+        .map(|rank| 1.0 / (rank as f64 + 2.0).log2())
+        .sum();
+    dcg / idcg
+}
+
+fn evaluate(model: &StaticModel, dir: &Path) -> f64 {
+    let (doc_ids, docs) = read_texts(&dir.join("corpus.jsonl"));
+    let (query_ids, queries) = read_texts(&dir.join("queries.jsonl"));
+    let mut qrels: HashMap<String, HashSet<String>> = HashMap::new();
+    for v in read_jsonl(&dir.join("qrels.jsonl")) {
+        qrels
+            .entry(field(&v, "query-id"))
+            .or_default()
+            .insert(field(&v, "corpus-id"));
+    }
+
+    let doc_emb: Vec<Vec<f32>> = model.encode(&docs).into_iter().map(normalized).collect();
+    let query_emb: Vec<Vec<f32>> = model.encode(&queries).into_iter().map(normalized).collect();
+
+    let mut total = 0.0;
+    let mut count = 0;
+    for (qid, q) in query_ids.iter().zip(&query_emb) {
+        let Some(relevant) = qrels.get(qid) else { continue };
+        let mut scored: Vec<(f32, usize)> = doc_emb
+            .iter()
+            .enumerate()
+            .filter(|(d, _)| &doc_ids[*d] != qid)
+            .map(|(d, e)| (q.iter().zip(e).map(|(a, b)| a * b).sum(), d))
+            .collect();
+        scored.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| doc_ids[b.1].cmp(&doc_ids[a.1])));
+        let ranked: Vec<usize> = scored.iter().take(K).map(|&(_, d)| d).collect();
+        total += ndcg_at_k(&ranked, &doc_ids, relevant);
+        count += 1;
+    }
+    assert!(count > 0, "no queries with relevance judgements in {}", dir.display());
+    total / count as f64
+}
+
+fn main() {
+    let args: HashMap<String, String> = env::args()
+        .skip(1)
+        .filter_map(|a| a.split_once('=').map(|(k, v)| (k.to_string(), v.to_string())))
+        .collect();
+    let data = args
+        .get("data")
+        .expect("usage: nanobeir data=<dir> [model=..] [quantize=none|int8]");
+    let model_id = args.get("model").map_or("minishlab/potion-base-8M", String::as_str);
+    let quantize = args.get("quantize").map_or("none", String::as_str);
+
+    let model = StaticModel::from_pretrained(model_id, None, None, None)
+        .unwrap_or_else(|e| panic!("failed to load {model_id}: {e}"));
+    let model = match quantize {
+        "none" => model,
+        "int8" => model.quantize_int8(),
+        other => panic!("unknown quantize mode {other:?}"),
+    };
+
+    let mut datasets: Vec<_> = fs::read_dir(data)
+        .unwrap_or_else(|e| panic!("failed to read {data}: {e}"))
+        .map(|e| e.expect("failed to read directory entry").path())
+        .filter(|p| p.join("corpus.jsonl").exists())
+        .collect();
+    datasets.sort();
+    assert!(!datasets.is_empty(), "no datasets found in {data}");
+
+    for dir in datasets {
+        let name = dir.file_name().unwrap().to_string_lossy();
+        println!("SCORE {name} {:.6}", evaluate(&model, &dir));
+    }
+}

benches/nanobeir_baseline.json

@@ -0,0 +1,122 @@
+{
+  "minishlab/potion-base-8M": {
+    "NanoArguAna": 0.407863,
+    "NanoClimateFEVER": 0.262234,
+    "NanoDBPedia": 0.520421,
+    "NanoFEVER": 0.530583,
+    "NanoFiQA2018": 0.279583,
+    "NanoHotpotQA": 0.525467,
+    "NanoMSMARCO": 0.420119,
+    "NanoNFCorpus": 0.230995,
+    "NanoNQ": 0.339209,
+    "NanoQuoraRetrieval": 0.935887,
+    "NanoSCIDOCS": 0.260832,
+    "NanoSciFact": 0.57738,
+    "NanoTouche2020": 0.473055
+  },
+  "minishlab/potion-base-8M:int8": {
+    "NanoArguAna": 0.407863,
+    "NanoClimateFEVER": 0.262234,
+    "NanoDBPedia": 0.518853,
+    "NanoFEVER": 0.526414,
+    "NanoFiQA2018": 0.280005,
+    "NanoHotpotQA": 0.529993,
+    "NanoMSMARCO": 0.420119,
+    "NanoNFCorpus": 0.232001,
+    "NanoNQ": 0.339209,
+    "NanoQuoraRetrieval": 0.935887,
+    "NanoSCIDOCS": 0.260624,
+    "NanoSciFact": 0.576504,
+    "NanoTouche2020": 0.471371
+  },
+  "minishlab/potion-retrieval-32M": {
+    "NanoArguAna": 0.48231,
+    "NanoClimateFEVER": 0.335031,
+    "NanoDBPedia": 0.564652,
+    "NanoFEVER": 0.67344,
+    "NanoFiQA2018": 0.365996,
+    "NanoHotpotQA": 0.602196,
+    "NanoMSMARCO": 0.410042,
+    "NanoNFCorpus": 0.334901,
+    "NanoNQ": 0.442824,
+    "NanoQuoraRetrieval": 0.897223,
+    "NanoSCIDOCS": 0.305196,
+    "NanoSciFact": 0.68316,
+    "NanoTouche2020": 0.527281
+  },
+  "minishlab/potion-retrieval-32M:int8": {
+    "NanoArguAna": 0.48231,
+    "NanoClimateFEVER": 0.335031,
+    "NanoDBPedia": 0.564459,
+    "NanoFEVER": 0.67344,
+    "NanoFiQA2018": 0.36592,
+    "NanoHotpotQA": 0.602572,
+    "NanoMSMARCO": 0.410042,
+    "NanoNFCorpus": 0.335081,
+    "NanoNQ": 0.442647,
+    "NanoQuoraRetrieval": 0.897223,
+    "NanoSCIDOCS": 0.305196,
+    "NanoSciFact": 0.683773,
+    "NanoTouche2020": 0.52625
+  },
+  "sentence-transformers/static-retrieval-mrl-en-v1": {
+    "NanoArguAna": 0.483885,
+    "NanoClimateFEVER": 0.316405,
+    "NanoDBPedia": 0.568139,
+    "NanoFEVER": 0.687252,
+    "NanoFiQA2018": 0.357502,
+    "NanoHotpotQA": 0.653112,
+    "NanoMSMARCO": 0.404068,
+    "NanoNFCorpus": 0.324124,
+    "NanoNQ": 0.453388,
+    "NanoQuoraRetrieval": 0.894883,
+    "NanoSCIDOCS": 0.265198,
+    "NanoSciFact": 0.609445,
+    "NanoTouche2020": 0.512113
+  },
+  "minishlab/potion-multilingual-128M": {
+    "NanoArguAna": 0.453841,
+    "NanoClimateFEVER": 0.266044,
+    "NanoDBPedia": 0.456969,
+    "NanoFEVER": 0.401784,
+    "NanoFiQA2018": 0.231971,
+    "NanoHotpotQA": 0.45701,
+    "NanoMSMARCO": 0.231029,
+    "NanoNFCorpus": 0.178024,
+    "NanoNQ": 0.213272,
+    "NanoQuoraRetrieval": 0.855965,
+    "NanoSCIDOCS": 0.146768,
+    "NanoSciFact": 0.524785,
+    "NanoTouche2020": 0.41905
+  },
+  "minishlab/potion-retrieval-32M:model2vec-int8": {
+    "NanoArguAna": 0.48249,
+    "NanoClimateFEVER": 0.332433,
+    "NanoDBPedia": 0.562414,
+    "NanoFEVER": 0.672588,
+    "NanoFiQA2018": 0.370841,
+    "NanoHotpotQA": 0.598369,
+    "NanoMSMARCO": 0.410042,
+    "NanoNFCorpus": 0.334431,
+    "NanoNQ": 0.443524,
+    "NanoQuoraRetrieval": 0.897401,
+    "NanoSCIDOCS": 0.305173,
+    "NanoSciFact": 0.683773,
+    "NanoTouche2020": 0.527582
+  },
+  "minishlab/potion-base-8M:model2vec-int8": {
+    "NanoArguAna": 0.413644,
+    "NanoClimateFEVER": 0.26009,
+    "NanoDBPedia": 0.520808,
+    "NanoFEVER": 0.535466,
+    "NanoFiQA2018": 0.290429,
+    "NanoHotpotQA": 0.521167,
+    "NanoMSMARCO": 0.420119,
+    "NanoNFCorpus": 0.231021,
+    "NanoNQ": 0.341983,
+    "NanoQuoraRetrieval": 0.935887,
+    "NanoSCIDOCS": 0.259793,
+    "NanoSciFact": 0.578403,
+    "NanoTouche2020": 0.475323
+  }
+}