File size: 7,644 Bytes
ab4f92b b46c110 4113449 ab4f92b 4113449 b46c110 4113449 ab4f92b b46c110 ab4f92b 4113449 ab4f92b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | """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()
|