"""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): @classmethod 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.") @unittest.skipUnless(shutil.which("bun"), "bun not installed") 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)) @unittest.skipUnless(shutil.which("bun"), "bun not installed") class TestWebExport(unittest.TestCase): """web/model.bin + web/autocomplete.ts (run `python export_web.py` first).""" @classmethod 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()