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]:
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()