chore: update benchmark in readme

Cargo.toml

@@ -67,3 +67,7 @@ harness = false
 [[bench]]
 name = "nanobeir"
 harness = false
+
+[[bench]]
+name = "python_compare"
+harness = false

README.md

@@ -209,18 +209,39 @@ A variety of pre-trained Model2Vec models are available on the [HuggingFace Hub
 
 ## Performance
 
-We compared the performance of the Rust implementation with the Python version of Model2Vec. The benchmark was run single-threaded on a CPU.
+We compared `model2vec-rs` with the Python version of Model2Vec by encoding the same texts with [`potion-base-8M`](https://huggingface.co/minishlab/potion-base-8M) in both.
 
-| Implementation | Throughput                                         |
-| -------------- | -------------------------------------------------- |
-| **Rust**       | 8000 samples/second |
-| **Python**     | 4650 samples/second |
+| Implementation | Single thread       | All cores             | Model load |
+| -------------- | ------------------- | --------------------- | ---------- |
+| **Rust**       | 671K texts/second   | 3.06M texts/second    | 15 ms      |
+| **Python**     | 40K texts/second    | 134K texts/second     | 312 ms     |
+| **Speedup**    | **16.6×**           | **22.9×**             | **~20×**   |
 
-The Rust version is roughly **1.7×** faster than the Python version.
+On a single thread, the Rust version is roughly **17×** faster than the Python version, and about **23×** faster when using all cores.
+
+To reproduce these numbers, run this from the repository root (it needs Rust and [`uv`](https://docs.astral.sh/uv/)):
+
+```sh
+uv run scripts/python_compare.py
+```
+
+The script downloads the corpus, builds [`benches/python_compare.rs`](benches/python_compare.rs), runs both implementations and prints the table above. See [`scripts/python_compare.py`](scripts/python_compare.py) for the options.
+
+Benchmark details:
+
+*   **Corpus:** 2,644 random articles from English Wikipedia: shard [`20231101.en/train-00017-of-00041.parquet`](https://huggingface.co/datasets/wikimedia/wikipedia/blob/b04c8d1ceb2f5cd4588862100d08de323dccfbaa/20231101.en/train-00017-of-00041.parquet) of [`wikimedia/wikipedia`](https://huggingface.co/datasets/wikimedia/wikipedia), shuffled with seed 0 and taken until they reach 1.1M words, joined by blank lines (7.0 MB, SHA-256 `c23440073f72900784a3f56a2318742ff055506b7a136f1e473edb4ef190d9e7`). The script rebuilds it byte for byte.
+*   **Texts:** each of the 64,738 non-empty lines (paragraphs and section headings, 17 words and 108 bytes on average) is one text. Every timed pass encodes all of them in a single `encode` call with the default `max_length` of 512.
+*   **Hardware:** Apple M3 Pro (12 cores), 18 GB RAM.
+*   **Software:** `model2vec-rs` 0.3.0 (release build, Rust 1.99.1); Python 3.14.3 with `model2vec` 0.10.0, `tokenizers` 0.23.3 and `numpy` 2.5.4.
+*   **Single thread:** `RAYON_NUM_THREADS=1` and `TOKENIZERS_PARALLELISM=false` for both; Python with `use_multiprocessing=False`.
+*   **All cores:** both with their defaults: Rust with Rayon parallelism; Python with tokenizer parallelism and `use_multiprocessing=True`, which spreads the texts over a process pool.
+*   **Method:** 5 alternating rounds of Rust and Python processes, 5 timed passes each after one warm-up pass; the table reports the median over the 25 passes.
+*   **Output:** the Rust and Python embeddings differ by at most 2.7e-7 per value (minimum cosine similarity 0.9999999).
+*   **Model load:** median time to load the model from the local Hugging Face cache; for Python this includes importing `model2vec`.
 
 ## Relation to Python `model2vec`
 
-*   **`model2vec-rs` (This Crate):** High-performance Rust engine for 1.7x faster **Model2Vec inference**.
+*   **`model2vec-rs` (This Crate):** High-performance Rust engine for up to 23x faster **Model2Vec inference**.
 *   **[`model2vec`](https://github.com/MinishLab/model2vec) (Python-based):** Handles model **distillation, training, fine-tuning**, and slower Python-based inference.
 
 ## License

benches/README.md

@@ -8,6 +8,8 @@
 - `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.
+- `python_compare.rs`: Rust side of the comparison with Python model2vec
+  in the main README. Normally run through `scripts/python_compare.py`.
 
 ## Comparing two revisions
 
@@ -43,6 +45,21 @@ For trustworthy numbers:
 4. Only claim a speedup for configurations marked `faster` whose output is
    `identical` (for changes that must not alter embeddings).
 
+## Comparing with Python model2vec
+
+```sh
+uv run scripts/python_compare.py
+```
+
+Reproduces the Performance table in the main README. The script rebuilds
+the pinned Wikipedia corpus under `target/python-compare/` (checked against
+a SHA-256 digest), uses each non-empty line as one text, and runs Rust and
+Python in alternating processes for `--threads 1,all`. Python uses its
+default process pool (`use_multiprocessing=True`) only for `all`; pass
+`--no-python-multiprocessing` to turn it off there too. It also reports how
+far the Rust embeddings are from the Python ones. Pass `--corpus` and
+`--model` to compare on other data or models.
+
 ## Int8 models
 
 To benchmark int8 storage, quantize models with Python model2vec and pass

benches/python_compare.rs

@@ -0,0 +1,60 @@
+// Rust side of the comparison with Python model2vec, driven by
+// `scripts/python_compare.py`.
+//
+//   cargo bench --bench python_compare -- corpus=wiki-en.txt model=minishlab/potion-base-8M \
+//       reps=5 out=embeddings.f32
+//
+// Every non-empty line of the corpus is one text, and each timed pass encodes
+// all of them with one `encode` call. Prints one `LOAD <seconds>` line, one
+// `TEXTS <n>` line and one `SAMPLE <seconds>` line per timed pass. With
+// `out=`, writes the embeddings of the warm-up pass as little-endian f32,
+// row-major.
+
+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 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 corpus = args
+        .get("corpus")
+        .expect("usage: python_compare corpus=<file> [model=..] [reps=N] [out=<file>]");
+    let model_id = args.get("model").map_or("minishlab/potion-base-8M", String::as_str);
+    let reps: usize = args
+        .get("reps")
+        .map_or(5, |r| r.parse().expect("reps must be an integer"));
+
+    let text = fs::read_to_string(corpus).unwrap_or_else(|e| panic!("failed to read {corpus}: {e}"));
+    let texts: Vec<String> = text
+        .lines()
+        .filter(|l| !l.trim().is_empty())
+        .map(String::from)
+        .collect();
+    assert!(!texts.is_empty(), "corpus {corpus} has no non-empty lines");
+
+    let start = Instant::now();
+    let model = StaticModel::from_pretrained(model_id, None, None, None)
+        .unwrap_or_else(|e| panic!("failed to load {model_id}: {e}"));
+    println!("LOAD {:.6}", start.elapsed().as_secs_f64());
+    println!("TEXTS {}", texts.len());
+
+    let warm = model.encode(&texts);
+    if let Some(out) = args.get("out") {
+        let bytes: Vec<u8> = warm.iter().flatten().flat_map(|x| x.to_le_bytes()).collect();
+        fs::write(out, bytes).unwrap_or_else(|e| panic!("failed to write {out}: {e}"));
+    }
+    drop(warm);
+
+    for _ in 0..reps {
+        let start = Instant::now();
+        black_box(model.encode(&texts));
+        println!("SAMPLE {:.9}", start.elapsed().as_secs_f64());
+    }
+}

scripts/python_compare.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["model2vec==0.10.0", "numpy", "pyarrow", "huggingface_hub"]
+# ///
+"""Throughput comparison of model2vec-rs with Python model2vec.
+
+Builds `benches/python_compare.rs` from this checkout, then encodes the same
+texts with both implementations in separate processes, alternating Rust and
+Python (R P, P R, ...) so slow drift hits both sides equally. Every non-empty
+line of the corpus is one text; each timed pass encodes all of them in one
+`encode` call with the default `max_length` of 512.
+
+The default corpus is the one used for the README numbers: 2,644 random
+articles from English Wikipedia (`wikimedia/wikipedia`, config `20231101.en`,
+shard 17, pinned revision), shuffled with seed 0 and taken until they reach
+1.1M words. It is downloaded once to `target/python-compare/wiki-en.txt` and
+checked against a SHA-256 digest.
+
+Usage (from the repository root):
+
+    uv run scripts/python_compare.py
+    uv run scripts/python_compare.py --threads 1 --rounds 3
+    uv run scripts/python_compare.py --corpus my-texts.txt --model minishlab/potion-retrieval-32M
+
+`--threads 1` sets `RAYON_NUM_THREADS=1` and `TOKENIZERS_PARALLELISM=false`
+for both sides; `all` leaves both at their default parallelism. Python runs
+with its default `use_multiprocessing=True` for `all`, which spreads large
+inputs over a process pool; `--no-python-multiprocessing` turns that off.
+Every other thread count runs Python with `use_multiprocessing=False`, since
+the process pool ignores the thread limit.
+The script also reports how far the Rust embeddings are from the Python ones.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import random
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+REPO = Path(__file__).resolve().parent.parent
+WORK = REPO / "target" / "python-compare"
+WIKI_REPO = "wikimedia/wikipedia"
+WIKI_REVISION = "b04c8d1ceb2f5cd4588862100d08de323dccfbaa"
+WIKI_SHARD = "20231101.en/train-00017-of-00041.parquet"
+WIKI_WORDS = 1_100_000
+WIKI_SHA256 = "c23440073f72900784a3f56a2318742ff055506b7a136f1e473edb4ef190d9e7"
+
+
+def sha256(path: Path) -> str:
+    return hashlib.sha256(path.read_bytes()).hexdigest()
+
+
+def ensure_corpus() -> Path:
+    path = WORK / "wiki-en.txt"
+    if path.exists() and sha256(path) == WIKI_SHA256:
+        return path
+    import pyarrow.parquet as pq
+    from huggingface_hub import hf_hub_download
+
+    print(f"downloading {WIKI_REPO}/{WIKI_SHARD}@{WIKI_REVISION[:12]}", file=sys.stderr)
+    shard = hf_hub_download(WIKI_REPO, WIKI_SHARD, repo_type="dataset", revision=WIKI_REVISION)
+    articles = pq.read_table(shard, columns=["text"]).column("text").to_pylist()
+    random.Random(0).shuffle(articles)
+    sample, words = [], 0
+    for article in articles:
+        sample.append(article)
+        words += len(article.split())
+        if words >= WIKI_WORDS:
+            break
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes("\n\n".join(sample).encode("utf-8"))
+    if sha256(path) != WIKI_SHA256:
+        sys.exit(f"corpus {path} does not match the expected SHA-256 {WIKI_SHA256}")
+    return path
+
+
+def read_texts(corpus: Path) -> list[str]:
+    """Split `corpus` like `str::lines` in Rust and keep the non-empty lines."""
+    lines = corpus.read_bytes().decode("utf-8").split("\n")
+    return [line.removesuffix("\r") for line in lines if line.strip()]
+
+
+def python_worker(corpus: Path, model_id: str, reps: int, multiprocessing: bool, out: Path | None) -> None:
+    start = time.perf_counter()
+    from model2vec import StaticModel
+
+    model = StaticModel.from_pretrained(model_id)
+    print(f"LOAD {time.perf_counter() - start:.6f}")
+    texts = read_texts(corpus)
+    print(f"TEXTS {len(texts)}")
+    warm = model.encode(texts, use_multiprocessing=multiprocessing)
+    if out:
+        warm.astype("<f4").tofile(out)
+    del warm
+    for _ in range(reps):
+        start = time.perf_counter()
+        model.encode(texts, use_multiprocessing=multiprocessing)
+        print(f"SAMPLE {time.perf_counter() - start:.9f}")
+
+
+def build() -> Path:
+    print("building benches/python_compare.rs", file=sys.stderr)
+    out = subprocess.run(
+        ["cargo", "bench", "--bench", "python_compare", "--no-run", "--message-format=json"],
+        cwd=REPO,
+        check=True,
+        text=True,
+        capture_output=True,
+    ).stdout
+    for line in out.splitlines():
+        msg = json.loads(line)
+        if msg.get("reason") == "compiler-artifact" and msg["target"]["name"] == "python_compare" and msg.get("executable"):
+            return Path(msg["executable"])
+    raise RuntimeError("could not find the python_compare executable")
+
+
+def thread_env(threads: str) -> dict[str, str]:
+    env = dict(os.environ)
+    if threads != "all":
+        env["RAYON_NUM_THREADS"] = threads
+        env["RAYON_RS_NUM_THREADS"] = threads
+        env["OMP_NUM_THREADS"] = threads
+        if threads == "1":
+            env["TOKENIZERS_PARALLELISM"] = "false"
+    return env
+
+
+def run_side(cmd: list[str], threads: str) -> tuple[float, int, list[float]]:
+    out = subprocess.run(cmd, env=thread_env(threads), check=True, text=True, capture_output=True).stdout
+    load, texts, samples = 0.0, 0, []
+    for line in out.splitlines():
+        key, _, value = line.partition(" ")
+        if key == "LOAD":
+            load = float(value)
+        elif key == "TEXTS":
+            texts = int(value)
+        elif key == "SAMPLE":
+            samples.append(float(value))
+    return load, texts, samples
+
+
+def compare_embeddings(rust: Path, python: Path, n: int) -> dict[str, float]:
+    import numpy as np
+
+    a = np.fromfile(rust, dtype="<f4").reshape(n, -1)
+    b = np.fromfile(python, dtype="<f4").reshape(n, -1)
+    norms = np.linalg.norm(a, axis=1) * np.linalg.norm(b, axis=1)
+    nonzero = norms > 0
+    cos = (a * b).sum(axis=1)[nonzero] / norms[nonzero]
+    return {
+        "dim": int(a.shape[1]),
+        "max_abs_diff": float(np.abs(a - b).max()),
+        "min_cosine": float(cos.min()),
+        "identical_rows": int((a == b).all(axis=1).sum()),
+    }
+
+
+def fmt_rate(x: float) -> str:
+    for unit, scale in (("M", 1e6), ("K", 1e3)):
+        if x >= scale:
+            return f"{x / scale:.2f}{unit}"
+    return f"{x:.0f}"
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", default="minishlab/potion-base-8M")
+    p.add_argument("--corpus", type=Path, help="text file, one text per line (default: the pinned Wikipedia sample)")
+    p.add_argument("--threads", default="1,all", help="comma-separated: 'all' and/or thread counts")
+    p.add_argument("--rounds", type=int, default=5, help="alternating process launches per side")
+    p.add_argument("--reps", type=int, default=5, help="timed passes per process launch")
+    p.add_argument("--no-python-multiprocessing", action="store_true", help="use_multiprocessing=False in Python for 'all' too")
+    p.add_argument("--python-multiprocessing", action="store_true", help=argparse.SUPPRESS)
+    p.add_argument("--out", type=Path, help="write raw results as JSON (default: target/python-compare/results-*.json)")
+    p.add_argument("--python-worker", action="store_true", help=argparse.SUPPRESS)
+    p.add_argument("--embeddings", type=Path, help=argparse.SUPPRESS)
+    a = p.parse_args()
+
+    if a.python_worker:
+        python_worker(a.corpus, a.model, a.reps, a.python_multiprocessing, a.embeddings)
+        return
+
+    corpus = a.corpus.resolve() if a.corpus else ensure_corpus()
+    binary = build()
+    WORK.mkdir(parents=True, exist_ok=True)
+    rust_base = [str(binary), f"corpus={corpus}", f"model={a.model}"]
+    python_base = [sys.executable, str(Path(__file__).resolve()), "--python-worker", "--corpus", str(corpus), "--model", a.model]
+    python_mp = not a.no_python_multiprocessing
+    rust_cmd = [*rust_base, f"reps={a.reps}"]
+
+    def python_cmd(threads: str) -> list[str]:
+        mp = ["--python-multiprocessing"] if python_mp and threads == "all" else []
+        return [*python_base, *mp, "--reps", str(a.reps)]
+
+    rust_emb, python_emb = WORK / "rust.f32", WORK / "python.f32"
+    _, n, _ = run_side([*rust_base, "reps=0", f"out={rust_emb}"], "all")
+    run_side([*python_base, "--reps", "0", "--embeddings", str(python_emb)], "all")
+    output = compare_embeddings(rust_emb, python_emb, n)
+
+    results = {}
+    started = time.time()
+    for threads in a.threads.split(","):
+        res = results[threads] = {"rust": {"load": [], "seconds": []}, "python": {"load": [], "seconds": []}}
+        for r in range(a.rounds):
+            print(f"threads={threads} round {r + 1}/{a.rounds}", file=sys.stderr)
+            order = ("rust", "python") if r % 2 == 0 else ("python", "rust")
+            for side in order:
+                load, _, samples = run_side(rust_cmd if side == "rust" else python_cmd(threads), threads)
+                res[side]["load"].append(load)
+                res[side]["seconds"].extend(samples)
+
+    print(f"\nmodel={a.model} corpus={corpus} texts={n} rounds={a.rounds} reps={a.reps}")
+    print(f"python use_multiprocessing={python_mp} for threads=all   elapsed {time.time() - started:.0f}s\n")
+    header = f"{'threads':>7} {'Rust texts/s':>14} {'Python texts/s':>15} {'speedup':>8} {'Rust load':>10} {'Python load':>12}"
+    print(header)
+    print("-" * len(header))
+    for threads, res in results.items():
+        rate = {side: statistics.median(n / s for s in res[side]["seconds"]) for side in res}
+        load = {side: statistics.median(res[side]["load"]) for side in res}
+        res["median_texts_per_second"], res["median_load_seconds"] = rate, load
+        print(
+            f"{threads:>7} {fmt_rate(rate['rust']):>14} {fmt_rate(rate['python']):>15} "
+            f"{rate['rust'] / rate['python']:>7.1f}x {load['rust'] * 1e3:>8.0f}ms {load['python'] * 1e3:>10.0f}ms"
+        )
+    print(
+        f"\nRust vs Python embeddings ({output['dim']} dims): max abs diff {output['max_abs_diff']:.2e}, "
+        f"min cosine {output['min_cosine']:.7f}, {output['identical_rows']}/{n} rows bit-identical"
+    )
+    print("Throughput is the median over all timed passes; Python load time includes importing model2vec.")
+
+    out = a.out or WORK / f"results-{time.strftime('%Y%m%d-%H%M%S')}.json"
+    meta = {"model": a.model, "corpus": str(corpus), "texts": n, "rounds": a.rounds, "reps": a.reps, "python_multiprocessing_for_all": python_mp}
+    out.write_text(json.dumps({"meta": meta, "output": output, "results": results}, indent=2))
+    print(f"raw results: {out}")
+
+
+if __name__ == "__main__":
+    main()