| """Write web/model.bin: the compiled model in compact, quantized form for the browser. |
| |
| The dense model.safetensors is ~93% zeros. Each MLP neuron of the compiled |
| transformer is fully described by its context tokens (which slots its gate |
| matches) and the logit changes its down-projection writes; the attention and |
| embedding weights are fixed by construction. This file stores exactly that, so |
| web/autocomplete.ts can evaluate only the neurons whose context matches: the |
| same computation as the transformer, without the zeros. |
| |
| Quantization (format v2): |
| - logit changes are always >= 0, stored as uint8 with one scale per context length |
| ("u8"), or as float16 ("f16"), or unquantized ("f32", used by the exactness test); |
| see compile_settings.json "export" |
| - the out-of-vocabulary prefix table keeps only prefixes of vocabulary words (the |
| only ones the policy ever looks up), with rates as uint8 / 255 |
| - context counts are omitted unless the shipped policy uses min_support |
| - contexts are sorted, so the runtime binary-searches them instead of building a map |
| |
| Layout (little-endian): |
| "GHST" | u32 version | u32 header bytes | header JSON | pad to 8 |
| sections (each padded to 8), offsets relative to the first section: |
| vocab utf-8, pieces joined by "\\n" |
| oov_keys utf-8, prefixes joined by "\\n"; oov_vals u8[n] |
| base f32[V] unigram neuron (always on) |
| ctx{k} k=1..order u16/u32[n_k * k] context ids, oldest first, sorted |
| count{k} u16[n_k] number of logit changes per context |
| dids{k} u16/u32[pairs_k] next-token ids |
| dvals{k} u8 / f16 / f32[pairs_k] logit changes (u8: value * scale_k) |
| support{k} f32[n_k] only if policy.min_support > 0 |
| |
| python export_web.py # "web" settings, "policy" and "export" in compile_settings.json |
| """ |
|
|
| import json |
| from dataclasses import asdict |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| import build_weights as bw |
| import suggest as sg |
|
|
| ROOT = Path(__file__).parent |
| OUT = ROOT / "web" / "model.bin" |
| MAGIC, VERSION = b"GHST", 2 |
|
|
|
|
| def load_policy() -> sg.Policy: |
| return sg.Policy(**json.loads(bw.SETTINGS_PATH.read_text())["policy"]) |
|
|
|
|
| def load_scheme() -> str: |
| return json.loads(bw.SETTINGS_PATH.read_text()).get("export", {}).get("deltas", "u8") |
|
|
|
|
| def quantize_oov(oov: dict[str, float], vocab: list[str]) -> dict[str, int]: |
| """Keep prefixes of vocabulary words only (others never have candidates), rates as u8.""" |
| prefixes = {p[:i] for p in vocab for i in range(1, len(p) + 1)} |
| out = {} |
| for p, r in oov.items(): |
| q = int(round(r * 255)) |
| if q > 0 and p in prefixes: |
| out[p] = q |
| return out |
|
|
|
|
| def quantize(ngram: bw.NGram, vocab: list[str], oov: dict[str, float], scheme: str) -> dict: |
| """Quantized tables; also patches `ngram` so the Python reference uses the same numbers.""" |
| tables, deq = {}, {} |
| for k in range(1, ngram.order + 1): |
| ctxs = sorted(ngram.tables[k]) |
| counts, dids, dvals = [], [], [] |
| for ctx in ctxs: |
| d = ngram.delta(ctx) |
| counts.append(len(d)) |
| dids.extend(d.keys()) |
| dvals.extend(d.values()) |
| vals = np.array(dvals, dtype=np.float64) |
| assert (vals >= 0).all(), "deltas are log(A + g*p) - log(g*p) >= 0" |
| if scheme == "u8": |
| scale = float(vals.max()) / 255 if len(vals) else 1.0 |
| q = np.round(vals / scale).clip(0, 255).astype(np.uint8) |
| back = q.astype(np.float64) * scale |
| elif scheme in ("f16", "f32"): |
| scale, q = 1.0, vals.astype(np.float16 if scheme == "f16" else np.float32) |
| back = q.astype(np.float64) |
| else: |
| raise ValueError(scheme) |
| tables[k] = {"ctxs": ctxs, "counts": counts, "dids": dids, "q": q, "scale": scale} |
| pos = 0 |
| for ctx, n in zip(ctxs, counts): |
| deq[ctx] = (np.array(dids[pos : pos + n], dtype=np.int64), back[pos : pos + n]) |
| pos += n |
| base = ngram.base_logits.astype(np.float32).astype(np.float64) |
| ngram.base_logits = base |
| ngram.delta_arrays = lambda ctx: deq[ctx] |
| oov_q = quantize_oov(oov, vocab) |
| return {"tables": tables, "oov_q": oov_q, "oov": {p: q / 255 for p, q in oov_q.items()}} |
|
|
|
|
| def export(vocab: list[str], ngram: bw.NGram, oov: dict, policy: sg.Policy, path: Path, scheme: str = "u8") -> dict: |
| id_dtype = np.uint16 if len(vocab) < 65536 else np.uint32 |
| qt = quantize(ngram, vocab, oov, scheme) |
| sections, orders = [], [] |
|
|
| def add(name: str, data: bytes) -> None: |
| sections.append((name, data)) |
|
|
| add("vocab", "\n".join(vocab).encode()) |
| add("oov_keys", "\n".join(qt["oov_q"]).encode()) |
| add("oov_vals", np.array(list(qt["oov_q"].values()), dtype=np.uint8).tobytes()) |
| add("base", ngram.base_logits.astype(np.float32).tobytes()) |
| for k, t in qt["tables"].items(): |
| add(f"ctx{k}", np.array(t["ctxs"], dtype=id_dtype).reshape(-1).tobytes()) |
| add(f"count{k}", np.array(t["counts"], dtype=np.uint16).tobytes()) |
| add(f"dids{k}", np.array(t["dids"], dtype=id_dtype).tobytes()) |
| add(f"dvals{k}", t["q"].tobytes()) |
| if policy.min_support > 0: |
| add(f"support{k}", np.array([ngram.support[c] for c in t["ctxs"]], dtype=np.float32).tobytes()) |
| orders.append({"k": k, "contexts": len(t["ctxs"]), "pairs": len(t["dids"]), "scale": t["scale"]}) |
|
|
| table, offset = [], 0 |
| for name, data in sections: |
| table.append({"name": name, "offset": offset, "bytes": len(data)}) |
| offset += -(-len(data) // 8) * 8 |
| header = { |
| "n_prev": ngram.s.n_prev, |
| "vocab_size": len(vocab), |
| "eos_id": bw.EOS_ID, |
| "unk_id": bw.UNK_ID, |
| "id_bytes": np.dtype(id_dtype).itemsize, |
| "deltas": scheme, |
| "orders": orders, |
| "policy": asdict(policy), |
| "stop_after": sorted(sg.STOP_AFTER), |
| "settings": asdict(ngram.s), |
| "sections": table, |
| } |
| head = json.dumps(header, separators=(",", ":")).encode() |
| path.parent.mkdir(exist_ok=True) |
| with open(path, "wb") as f: |
| f.write(MAGIC + np.array([VERSION, len(head)], dtype="<u4").tobytes() + head) |
| f.write(b"\0" * (-(12 + len(head)) % 8)) |
| for _, data in sections: |
| f.write(data + b"\0" * (-len(data) % 8)) |
| return { |
| "bytes": path.stat().st_size, |
| "contexts": sum(o["contexts"] for o in orders), |
| "pairs": sum(o["pairs"] for o in orders), |
| "oov_prefixes": len(qt["oov_q"]), |
| "oov": qt["oov"], |
| } |
|
|
|
|
| def build(settings: bw.Settings, sources: dict, policy: sg.Policy, path: Path = OUT, scheme: str = "u8"): |
| """Compile, quantize and write; returns (info, Suggester using the exported numbers).""" |
| vocab, ngram, stats = bw.compile_tables(settings, sources) |
| oov = sg.oov_prefix_rates(settings, sources, vocab) |
| info = export(vocab, ngram, oov, policy, path, scheme) |
| sugg = sg.Suggester(vocab, ngram, info.pop("oov")) |
| return {**stats, **info}, sugg |
|
|
|
|
| def main() -> None: |
| info, _ = build(bw.load_settings("web"), bw.load_sources(), load_policy(), OUT, load_scheme()) |
| print(json.dumps({**info, "mb": round(info["bytes"] / 1e6, 1), "out": str(OUT)}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
