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 | """Load the saved files exactly as a Hugging Face user would and check them. | |
| Run `python build_weights.py` first. | |
| """ | |
| import json | |
| import random | |
| import shutil | |
| import subprocess | |
| import unittest | |
| from pathlib import Path | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline | |
| import build_weights | |
| MODEL_DIR = Path(__file__).parent | |
| # prompt -> greedy completion (new text only) | |
| CASES = { | |
| "Once upon a": "time.", | |
| "The quick brown": "fox jumps over the lazy dog.", | |
| "Thank you for your": "help.", | |
| "I hope this email": "finds you well.", | |
| "Can I have the": "menu, please?", | |
| "Looking forward to": "working with you.", | |
| "Actions speak": "louder than words.", | |
| "Where is the": "bathroom?", | |
| "Happy birthday": ".", | |
| "I'm so proud": "of you.", | |
| # coding prompts and replies (data/*.txt) | |
| "In Python,": "you can use the following code:", | |
| "You can use the": "following code:", | |
| "How to convert a": "string to a file in python", | |
| } | |
| class TestModel(unittest.TestCase): | |
| def setUpClass(cls): | |
| cls.tok = AutoTokenizer.from_pretrained(MODEL_DIR) | |
| cls.model = AutoModelForCausalLM.from_pretrained(MODEL_DIR).eval() | |
| _, cls.vocab, cls.ngram, _, _, cls.info = build_weights.compile_model() | |
| def complete(self, prompt: str) -> str: | |
| ids = self.tok(prompt, return_tensors="pt") | |
| out = self.model.generate(**ids, max_new_tokens=16) | |
| return self.tok.decode(out[0, ids["input_ids"].shape[1] :], skip_special_tokens=True).strip() | |
| def test_parameter_count(self): | |
| n = sum(p.numel() for p in self.model.parameters()) | |
| self.assertEqual(n, self.info["params"]) | |
| self.assertLessEqual(n, 1_000_000_000) | |
| def test_greedy_completions(self): | |
| for prompt, expected in CASES.items(): | |
| with self.subTest(prompt=prompt): | |
| self.assertEqual(self.complete(prompt), expected) | |
| def test_matches_ngram_reference(self): | |
| """The transformer's next-token distribution equals the compiled n-gram's.""" | |
| index = {p: i for i, p in enumerate(self.vocab)} | |
| rng = random.Random(0) | |
| phrases = build_weights.load_phrases() | |
| cases = [] | |
| for toks in rng.sample(phrases, 200): | |
| ids = [0] + [index[p] for p in toks] | |
| cases.append(ids[: rng.randint(1, len(ids))]) | |
| for _ in range(50): # unseen word salad, up to 80 tokens | |
| cases.append([0] + [rng.randrange(2, len(self.vocab)) for _ in range(rng.randint(1, 80))]) | |
| with torch.no_grad(): | |
| for ids in cases: | |
| got = self.model(torch.tensor([ids])).logits[0, -1].double().log_softmax(-1) | |
| ref = self.ngram.reference_logits(ids).log_softmax(-1) | |
| self.assertLess((got - ref).abs().max().item(), 1e-3) # observed ~3e-5 | |
| def test_unknown_words_back_off(self): | |
| # "zebra" is out of vocabulary; the model backs off to the shorter context. | |
| self.assertEqual(self.complete("Zebra, thank you for your"), "help.") | |
| def test_pipeline(self): | |
| gen = pipeline("text-generation", model=str(MODEL_DIR)) | |
| self.assertEqual(gen("Once upon a")[0]["generated_text"], "Once upon a time.") | |
| def test_typescript_matches(self): | |
| prompts = ["once upon a", "how are", "zebra", "i hope this email"] | |
| out = subprocess.run( | |
| ["bun", "autocomplete.ts", "--json", *prompts], | |
| cwd=MODEL_DIR, capture_output=True, text=True, check=True, | |
| ) | |
| results = json.loads(out.stdout) | |
| with torch.no_grad(): | |
| for prompt, res in zip(prompts, results): | |
| with self.subTest(prompt=prompt): | |
| ids = self.tok(prompt, return_tensors="pt") | |
| logp = self.model(**ids).logits[0, -1].log_softmax(-1) | |
| top = logp.topk(3) | |
| self.assertEqual([t["id"] for t in res["top"]], top.indices.tolist()) | |
| for t, lp in zip(res["top"], top.values.tolist()): | |
| self.assertAlmostEqual(t["logprob"], lp, places=3) | |
| self.assertEqual(res["completion"].strip(), self.complete(prompt)) | |
| class TestWebExport(unittest.TestCase): | |
| """web/model.bin + web/autocomplete.ts (run `python export_web.py` first).""" | |
| def setUpClass(cls): | |
| import export_web | |
| import suggest | |
| cls.export_web, cls.suggest = export_web, suggest | |
| cls.settings, cls.sources = build_weights.load_settings("web"), build_weights.load_sources() | |
| cls.policy = export_web.load_policy() | |
| vocab, ngram, _ = build_weights.compile_tables(cls.settings, cls.sources) | |
| oov = suggest.oov_prefix_rates(cls.settings, cls.sources, vocab) | |
| cls.suggester = suggest.Suggester(vocab, ngram, oov) | |
| def cli(self, *args: str) -> list: | |
| out = subprocess.run(["bun", "web/cli.ts", *args], cwd=MODEL_DIR, capture_output=True, text=True, check=True) | |
| return json.loads(out.stdout) | |
| def test_size(self): | |
| self.assertLessEqual((MODEL_DIR / "web" / "model.bin").stat().st_size, 50_000_000) | |
| def test_matches_python_policy(self): | |
| """Same suggestion chains (text, probability, support) as suggest.py on random prefixes.""" | |
| rng = random.Random(1) | |
| lines = [ln for ln in (MODEL_DIR / "data" / "prompts.txt").read_text().splitlines() if len(ln) > 2] | |
| prefixes = [t[: rng.randint(1, len(t) - 1)] for t in rng.sample(lines, 300)] | |
| for p, res in zip(prefixes, self.cli("--json", *prefixes)): | |
| with self.subTest(prefix=p): | |
| py = self.suggester.steps(p, self.policy.max_words) | |
| self.assertEqual([s.text for s in py], [s["text"] for s in res["steps"]]) | |
| for a, b in zip(py, res["steps"]): | |
| self.assertAlmostEqual(a.prob, b["prob"], places=3) | |
| self.assertAlmostEqual(a.support, b["support"], places=3) | |
| if all(abs(s.prob - self.policy.tau) > 1e-3 for s in py): # skip float32 ties at the threshold | |
| self.assertEqual(self.suggest.Suggester.apply(py, self.policy), res["ghost"]) | |
| def test_runtime_matches_dense_transformer(self): | |
| """The sparse runtime, run on the dense model's own tables, reproduces PyTorch.""" | |
| import tempfile | |
| s = build_weights.load_settings("hf") | |
| vocab, ngram, _ = build_weights.compile_tables(s, self.sources) | |
| oov = self.suggest.oov_prefix_rates(s, self.sources, vocab) | |
| texts = ["Once upon a", "write a python function to", "zebra", "how do i", "thank you for your"] | |
| tok = AutoTokenizer.from_pretrained(MODEL_DIR) | |
| model = AutoModelForCausalLM.from_pretrained(MODEL_DIR).eval() | |
| with tempfile.TemporaryDirectory() as tmp: | |
| path = Path(tmp) / "model.bin" | |
| self.export_web.export(vocab, ngram, oov, self.policy, path) | |
| results = self.cli("--model", str(path), "--dist", *texts) | |
| with torch.no_grad(): | |
| for t, res in zip(texts, results): | |
| with self.subTest(text=t): | |
| lp = model(**tok(t, return_tensors="pt")).logits[0, -1].double().log_softmax(-1) | |
| self.assertEqual([i for i, _ in res["top"]], lp.topk(5).indices.tolist()) | |
| for i, v in res["top"]: | |
| self.assertAlmostEqual(v, lp[i].item(), places=3) | |
| if __name__ == "__main__": | |
| unittest.main() | |