feat: add faster pooling

.gitignore

@@ -1,3 +1,4 @@
+__pycache__/*
 # Generated by Cargo
 # will have compiled files and executables
 debug/
@@ -18,4 +19,4 @@ target/
 
 local
 .DS_Store
-data
+data

Cargo.toml

@@ -56,3 +56,7 @@ required-features = ["cli"]
 [[bench]]
 name = "encode"
 harness = false
+
+[[bench]]
+name = "throughput"
+harness = false

benches/README.md

@@ -0,0 +1,40 @@
+# Benchmarks
+
+- `encode.rs`: latency percentiles of `StaticModel::encode` for one revision.
+- `throughput.rs`: throughput of one (model, workload, mode) configuration.
+  Uses only the public API, so it builds against any revision. Normally run
+  through `scripts/bench_compare.py`, not directly.
+
+## Comparing two revisions
+
+```sh
+python3 scripts/bench_compare.py                      # main vs the working tree
+python3 scripts/bench_compare.py --baseline v0.3.0 --candidate my-branch
+```
+
+The script checks the baseline out into a git worktree under
+`target/bench-compare/`, copies this `throughput.rs` into it, builds each
+commit (and the working tree) into its own target directory, refuses to run
+if both sides produce the same binary, and runs them alternately
+(A B, B A, ...) so drift affects both equally. Per configuration it reports
+median throughput, the paired speedup with a 95% bootstrap confidence
+interval, and whether both sides produced bit-identical embeddings. Raw
+samples are written to `target/bench-compare/results-*.json`.
+
+Configurations are the product of `--models`, `--workloads` (`queries`:
+8-word texts, `docs`: 256-word texts from `big.txt`), `--modes` (`batch`:
+`encode`, `single`: `encode_single` in a loop) and `--threads` (`all`, or a
+count; `1` also disables tokenizer parallelism). The full default run takes
+about 35 minutes, most of it on the two Unigram models, whose tokenization
+is far slower than pooling; pass `--models` to restrict it.
+
+For trustworthy numbers:
+
+1. Run an A/A comparison first (`--baseline main --candidate main`). Its
+   speedups show the noise floor of the machine; differences of that size
+   in a real comparison mean nothing.
+2. Use the default 10 rounds or more; fewer than 5 rounds are not flagged.
+3. Keep the machine idle and on power. `--threads 1` is the least noisy
+   setting and isolates per-text cost from scheduling.
+4. Only claim a speedup for configurations marked `faster` whose output is
+   `identical` (for changes that must not alter embeddings).

benches/throughput.rs

@@ -0,0 +1,85 @@
+// Throughput benchmark for one (model, workload, mode) configuration, driven
+// by `scripts/bench_compare.py`. Uses only the public API so the same file
+// builds against any revision.
+//
+//   cargo bench --bench throughput -- corpus=big.txt model=minishlab/potion-base-8M \
+//       workload=queries mode=batch reps=5
+//
+// Prints one `CONFIG` line, one `CHECKSUM` line (a hash of the bits of every
+// output value, to confirm two revisions compute the same embeddings), and
+// one `SAMPLE <seconds>` line per timed pass over the workload.
+
+use std::collections::HashMap;
+use std::env;
+use std::fs;
+use std::hint::black_box;
+use std::time::Instant;
+
+use model2vec_rs::model::StaticModel;
+
+fn checksum(embeddings: &[Vec<f32>]) -> u64 {
+    let mut h: u64 = 0xcbf29ce484222325;
+    for v in embeddings {
+        h = (h ^ v.len() as u64).wrapping_mul(0x100000001b3);
+        for x in v {
+            h = (h ^ x.to_bits() as u64).wrapping_mul(0x100000001b3);
+        }
+    }
+    h
+}
+
+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 arg = |k: &str, default: &str| args.get(k).cloned().unwrap_or_else(|| default.to_string());
+    let corpus = args
+        .get("corpus")
+        .expect("usage: throughput corpus=<file> [model=..] [workload=queries|docs] [mode=batch|single] [reps=N]");
+    let model_id = arg("model", "minishlab/potion-base-8M");
+    let workload = arg("workload", "queries");
+    let mode = arg("mode", "batch");
+    let reps: usize = arg("reps", "5").parse().expect("reps must be an integer");
+
+    let (words_per_text, texts_per_pass) = match (workload.as_str(), mode.as_str()) {
+        ("queries", "batch") => (8, 100_000),
+        ("queries", "single") => (8, 20_000),
+        ("docs", "batch") => (256, 4_000),
+        ("docs", "single") => (256, 1_000),
+        other => panic!("unknown workload/mode {other:?}"),
+    };
+    let text = fs::read_to_string(corpus).unwrap_or_else(|e| panic!("failed to read {corpus}: {e}"));
+    let words: Vec<&str> = text.split_whitespace().collect();
+    let texts: Vec<String> = words
+        .chunks(words_per_text)
+        .cycle()
+        .take(texts_per_pass)
+        .map(|c| c.join(" "))
+        .collect();
+
+    let model = StaticModel::from_pretrained(&model_id, None, None, None)
+        .unwrap_or_else(|e| panic!("failed to load {model_id}: {e}"));
+
+    let run = || -> Vec<Vec<f32>> {
+        match mode.as_str() {
+            "batch" => model.encode(&texts),
+            _ => texts.iter().map(|t| model.encode_single(t)).collect(),
+        }
+    };
+
+    let warm = run();
+    let dim = warm.first().map_or(0, Vec::len);
+    println!(
+        "CONFIG model={model_id} workload={workload} mode={mode} texts={} dim={dim}",
+        texts.len()
+    );
+    println!("CHECKSUM {:016x}", checksum(&warm));
+    drop(warm);
+
+    for _ in 0..reps {
+        let start = Instant::now();
+        black_box(run());
+        println!("SAMPLE {:.9}", start.elapsed().as_secs_f64());
+    }
+}