diff --git a/Cargo.toml b/Cargo.toml index 729aaf8..03a413c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,3 +67,7 @@ harness = false [[bench]] name = "nanobeir" harness = false + +[[bench]] +name = "python_compare" +harness = false diff --git a/README.md b/README.md index cd58fe0..2bf07d0 100644 --- a/README.md +++ b/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 diff --git a/benches/README.md b/benches/README.md index fc11304..d0188d9 100644 --- a/benches/README.md +++ b/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 diff --git a/benches/python_compare.rs b/benches/python_compare.rs new file mode 100644 index 0000000..21cae39 --- /dev/null +++ b/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 ` line, one +// `TEXTS ` line and one `SAMPLE ` 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 = 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= [model=..] [reps=N] [out=]"); + 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 = 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 = 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()); + } +} diff --git a/scripts/python_compare.py b/scripts/python_compare.py new file mode 100755 index 0000000..c732ef3 --- /dev/null +++ b/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(" 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=" 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()