File size: 9,356 Bytes
c160abb | 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 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """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 of the dense model (new text only). Regression snapshots:
# plain greedy decoding, without the confidence policy the apps use.
CASES = {
"You can use the": "following code:",
"can you explain the": "difference between the two lists in the code",
"write a react": "native",
"Where is the": "error.",
"Thank you for your": "help with a single line",
}
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.get(p, 1) 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"), self.complete("thank you for your"))
def test_pipeline(self):
gen = pipeline("text-generation", model=str(MODEL_DIR))
# transformers 4.x pipelines sample by default and don't prepend BOS unless asked; 5.x do neither.
out = gen("You can use the", add_special_tokens=True, do_sample=False)[0]["generated_text"]
self.assertEqual(out, "You can use the following code:")
@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))
class TestNoSecrets(unittest.TestCase):
"""People paste API keys into prompts; none may reach the data files or a vocabulary."""
def test_data_and_vocabularies(self):
import sys
sys.path.insert(0, str(MODEL_DIR / "scripts"))
from fetch_coding_prompts import looks_secret
for path in sorted((MODEL_DIR / "data").glob("*.txt")):
hits = [n for n, line in enumerate(path.read_text().splitlines(), 1) if looks_secret(line)]
self.assertEqual(hits, [], f"{path.name} lines")
vocabs = {"tokenizer.json": list(json.loads((MODEL_DIR / "tokenizer.json").read_text())["model"]["vocab"])}
raw = (MODEL_DIR / "web" / "model.bin").read_bytes()
header = json.loads(raw[12 : 12 + int.from_bytes(raw[8:12], "little")])
start = 12 + int.from_bytes(raw[8:12], "little")
start += -start % 8
sec = next(s for s in header["sections"] if s["name"] == "vocab")
vocabs["web/model.bin"] = raw[start + sec["offset"] : start + sec["offset"] + sec["bytes"]].decode().split("\n")
for name, words in vocabs.items():
self.assertEqual([w for w in words if looks_secret(w)], [], name)
@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 tempfile
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()
# Rebuild the export; the Python Suggester uses the same quantized numbers.
with tempfile.TemporaryDirectory() as tmp:
path = Path(tmp) / "model.bin"
_, cls.suggester = export_web.build(cls.settings, cls.sources, cls.policy, path, export_web.load_scheme())
cls.rebuilt = path.read_bytes()
def test_committed_export_is_current(self):
"""web/model.bin is exactly what export_web.py produces from the current sources and settings."""
self.assertEqual((MODEL_DIR / "web" / "model.bin").read_bytes(), self.rebuilt)
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)
if self.policy.min_support > 0: # supports are only exported when the policy uses them
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 (unquantized) 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, scheme="f32")
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()
|