Confidence-gated ghost text for coding prompts: Kneser-Ney compile, 4.5 MB browser export + zero-dep runtime, held-out precision/keystrokes-saved benchmarks
4113449 verified | """Measure compile cost, inference speed, exactness and held-out quality. | |
| Writes docs/benchmarks.json. Run from the repo root after build_weights.py and export_web.py: | |
| python scripts/benchmark.py | |
| """ | |
| import hashlib | |
| import json | |
| import os | |
| import platform | |
| import random | |
| import resource | |
| import shutil | |
| import statistics | |
| import subprocess | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import torch | |
| ROOT = Path(__file__).resolve().parent.parent | |
| sys.path.insert(0, str(ROOT)) | |
| sys.path.insert(0, str(ROOT / "scripts")) | |
| import build_weights as bw # noqa: E402 | |
| import export_web # noqa: E402 | |
| import suggest as sg # noqa: E402 | |
| import sweep # noqa: E402 | |
| PROMPTS = [ | |
| "Once upon a", "I hope this email", "The quick brown", "Thank you for", "Looking forward to", | |
| "In Python,", "How to convert a", "Write a python function to", "You can use the", "Here is an example of", | |
| ] | |
| def sha256(path: Path) -> str: | |
| h = hashlib.sha256() | |
| with open(path, "rb") as f: | |
| for chunk in iter(lambda: f.read(1 << 24), b""): | |
| h.update(chunk) | |
| return h.hexdigest() | |
| def hardware() -> dict: | |
| cpu = "" | |
| if sys.platform == "darwin": | |
| cpu = subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip() | |
| cpu = cpu or f"{platform.system()} {platform.machine()}" | |
| try: | |
| ram = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") | |
| except (ValueError, OSError): | |
| ram = 0 | |
| return {"cpu": cpu, "cores": os.cpu_count(), "ram_gb": round(ram / 2**30), "gpu": "none used", | |
| "python": platform.python_version(), "torch": torch.__version__} | |
| def timed_runs(cmd: list[str], out: Path, n: int = 2) -> dict: | |
| times, hashes = [], [] | |
| for _ in range(n): | |
| t0 = time.perf_counter() | |
| subprocess.run(cmd, cwd=ROOT, check=True, capture_output=True) | |
| times.append(time.perf_counter() - t0) | |
| hashes.append(sha256(out)) | |
| return { | |
| "seconds": round(min(times), 1), | |
| "runs": [round(t, 1) for t in times], | |
| "byte_identical_rebuild": len(set(hashes)) == 1, | |
| "sha256": hashes[0], | |
| "mb": round(out.stat().st_size / 1e6, 2), | |
| } | |
| def compile_dense() -> dict: | |
| runs = timed_runs([sys.executable, "build_weights.py"], ROOT / "model.safetensors") | |
| peak = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss | |
| t0 = time.perf_counter() | |
| sources = bw.load_sources() | |
| t1 = time.perf_counter() | |
| s = bw.load_settings("hf") | |
| vocab, ngram, stats = bw.compile_tables(s, sources) | |
| t2 = time.perf_counter() | |
| state, _, info = bw.build_dense(vocab, ngram, s) | |
| t3 = time.perf_counter() | |
| return { | |
| **runs, | |
| "peak_memory_gb": round(peak / 2**30 if sys.platform == "darwin" else peak / 2**20, 1), | |
| "tokenize_seconds": round(t1 - t0, 2), | |
| "count_seconds": round(t2 - t1, 2), | |
| "build_tensors_seconds": round(t3 - t2, 2), | |
| "source_lines": sum(len(v) for v in sources.values()), | |
| "params": info["params"], | |
| "nonzero_params": sum(int((t != 0).sum()) for t in state.values()), | |
| **stats, | |
| } | |
| def compile_web() -> dict: | |
| runs = timed_runs([sys.executable, "export_web.py"], ROOT / "web" / "model.bin") | |
| header_len = int.from_bytes((ROOT / "web" / "model.bin").read_bytes()[8:12], "little") | |
| header = json.loads((ROOT / "web" / "model.bin").read_bytes()[12 : 12 + header_len]) | |
| return {**runs, "contexts": sum(o["contexts"] for o in header["orders"]), "vocab": len(header["vocab"])} | |
| def torch_inference() -> dict: | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| t0 = time.perf_counter() | |
| tok = AutoTokenizer.from_pretrained(ROOT) | |
| model = AutoModelForCausalLM.from_pretrained(ROOT).eval() | |
| load = time.perf_counter() - t0 | |
| with torch.no_grad(): | |
| ids = tok("Write a python function to", return_tensors="pt") | |
| model(**ids) # warm-up | |
| t0 = time.perf_counter() | |
| for _ in range(20): | |
| model(**ids) | |
| forward_ms = (time.perf_counter() - t0) / 20 * 1000 | |
| new_tokens, t0 = 0, time.perf_counter() | |
| for p in PROMPTS: | |
| ids = tok(p, return_tensors="pt") | |
| out = model.generate(**ids, max_new_tokens=16) | |
| new_tokens += out.shape[1] - ids["input_ids"].shape[1] | |
| gen = time.perf_counter() - t0 | |
| return {"load_seconds": round(load, 1), "next_word_ms": round(forward_ms, 1), | |
| "generate_tokens_per_second": round(new_tokens / gen, 1), "threads": torch.get_num_threads()} | |
| def web_latency(texts: list[str]) -> dict | None: | |
| """Per-keystroke suggestion time of web/autocomplete.ts on held-out prefixes (Bun).""" | |
| if not shutil.which("bun"): | |
| return None | |
| rng = random.Random(0) | |
| prefixes = [t[: rng.randint(1, len(t) - 1)] for t in rng.sample(texts, 500)] | |
| out = json.loads(subprocess.run(["bun", "web/cli.ts", "--json", *prefixes], cwd=ROOT, check=True, | |
| capture_output=True, text=True).stdout) | |
| ms = sorted(r["ms"] for r in out[50:]) # skip JIT warm-up | |
| return {"runtime": "web/autocomplete.ts on Bun", "median_ms": round(statistics.median(ms), 2), | |
| "p95_ms": round(ms[int(0.95 * len(ms))], 2), "max_ms": round(ms[-1], 2)} | |
| def exactness() -> float: | |
| from transformers import AutoModelForCausalLM | |
| _, vocab, ngram, _, _, _ = bw.compile_model() | |
| model = AutoModelForCausalLM.from_pretrained(ROOT).eval() | |
| rng, worst = random.Random(0), 0.0 | |
| with torch.no_grad(): | |
| for _ in range(50): | |
| ids = [0] + [rng.randrange(2, len(vocab)) for _ in range(rng.randint(1, 80))] | |
| got = model(torch.tensor([ids])).logits[0, -1].double().log_softmax(-1) | |
| worst = max(worst, (got - ngram.reference_logits(ids).log_softmax(-1)).abs().max().item()) | |
| return worst | |
| def heldout_typing(name: str, texts: list[str], src: dict) -> dict: | |
| """Simulated typing on held-out prompts with the shipped policy, compiled without them.""" | |
| s = bw.load_settings(name) | |
| vocab, ngram, _ = bw.compile_tables(s, src) | |
| sugg = sg.Suggester(vocab, ngram, sg.oov_prefix_rates(s, src, vocab)) | |
| recs = sweep.record(sugg, texts) | |
| policy = export_web.load_policy() | |
| curve = {f"tau={t}": sweep.evaluate(recs, sg.Policy(tau=t, min_support=policy.min_support)) | |
| for t in [policy.tau, 0.8, 0.9, 0.95]} | |
| return {"policy": policy.__dict__, "heldout_prompts": len(texts), "characters": sum(map(len, texts)), **curve} | |
| def heldout_next_token(texts_src: dict) -> dict: | |
| """Top-1/top-3 next-token accuracy of the dense model's tables on held-out prompts.""" | |
| s = bw.load_settings("hf") | |
| _, test = bw.heldout_split(bw.load_sources()["prompts"]) | |
| vocab, ngram, _ = bw.compile_tables(s, texts_src) | |
| index = {p: i for i, p in enumerate(vocab)} | |
| top1 = top3 = total = oov = 0 | |
| for toks in test: | |
| ids = [0] + [index.get(p, 1) for p in toks] + [0] | |
| for t in range(len(ids) - 1): | |
| target = ids[t + 1] | |
| total += 1 | |
| if target == 1: | |
| oov += 1 | |
| continue | |
| ranked = ngram.reference_logits(ids[: t + 1]).topk(3).indices.tolist() | |
| top1 += ranked[0] == target | |
| top3 += target in ranked | |
| return {"heldout_lines": len(test), "heldout_tokens": total, "next_token_top1": round(top1 / total, 3), | |
| "next_token_top3": round(top3 / total, 3), "oov_rate": round(oov / total, 3)} | |
| def main() -> None: | |
| results = {"hardware": hardware()} | |
| print("compiling dense model (2 runs)...", flush=True) | |
| results["compile_dense"] = compile_dense() | |
| print("compiling web export (2 runs)...", flush=True) | |
| results["compile_web"] = compile_web() | |
| print("inference...", flush=True) | |
| results["pytorch"] = torch_inference() | |
| texts, src = sweep.heldout_texts(), sweep.train_sources() | |
| results["web_keystroke_latency"] = web_latency(texts) | |
| results["max_logprob_error_vs_ngram"] = exactness() | |
| print("held-out quality...", flush=True) | |
| results["heldout_typing_web"] = heldout_typing("web", texts, src) | |
| results["heldout_typing_dense_tables"] = heldout_typing("hf", texts, src) | |
| results["heldout_next_token_dense"] = heldout_next_token(src) | |
| (ROOT / "docs" / "benchmarks.json").write_text(json.dumps(results, indent=2) + "\n") | |
| print(json.dumps(results, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |