File size: 4,335 Bytes
888eee4 | 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 | """Generate the Bend parity fixture under tests/fixtures/parity/.
A deterministic, randomly-initialised tiny Groot model plus one request blob
and the PyTorch reference outputs. `bend/groot/parity.mjs` replays the request
through the Bend engine and compares.
python scripts/make_parity_fixture.py [out_dir]
Produces:
ckpt/ save_pretrained tiny model (groot_config.json, encoder/, model.safetensors)
model.gbw flat Bend weight file (groot1.export.write_gbw)
req.bin u32 request blob (groot1.bendio.RequestBlob)
ref.npy float32 [probs(nopt) ; conf_logits(nq) ; logits(nopt)]
meta.json {data_offset, n_chunks, n_questions, n_options}
"""
from __future__ import annotations
import json
import sys
from pathlib import Path
import numpy as np
import torch
from groot1.bendio import RequestBlob, load_gbw, data_offset
from groot1.config import GrootConfig
from groot1.export import write_gbw
from groot1.model import GrootDecisionModel, collate
from groot1.serialize import EncodedQuestion, EncodedState
CLS, SEP, MASK, PAD = 2, 1, 4, 0
def make_question(qtype: int, head: list[int], state_head: list[int], options: list[list[int]]) -> EncodedQuestion:
"""Same layout as groot1.serialize.Serializer._question (shared option base, option blocks)."""
ids = [CLS, *head, SEP] + ([*state_head, SEP] if state_head else [])
pos = list(range(len(ids)))
blocks = [0] * len(ids)
base = len(ids)
markers = []
for k, opt in enumerate(options):
markers.append(len(ids))
ids += [MASK, *opt]
blocks += [k + 1] * (1 + len(opt))
pos += list(range(base, base + 1 + len(opt)))
return EncodedQuestion(name="q", qtype=qtype, ids=ids, position_ids=pos, markers=markers, blocks=blocks)
def build_request(rng: np.random.Generator, vocab: int, early_state_tokens: int = 8):
def ids(n):
return rng.integers(5, vocab, size=n).tolist()
chunks = [[CLS, *ids(n), SEP] for n in (10, 12)]
state_head = [t for c in chunks for t in c[1:-1]][:early_state_tokens]
# one of each primitive, options of different lengths
questions = []
for qtype, k in ((0, 3), (1, 4), (2, 2)):
questions.append(make_question(qtype, ids(5), state_head, [ids(2 + (j % 3)) for j in range(k)]))
return chunks, questions
def main() -> None:
out = Path(sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/parity")
ckpt = out / "ckpt"
out.mkdir(parents=True, exist_ok=True)
torch.manual_seed(0)
rng = np.random.default_rng(0)
cfg = GrootConfig(
name="groot-parity",
hidden_size=32,
chunk_tokens=32,
max_chunks=8,
tree_levels=4,
fusion_layers=2,
fusion_heads=4,
fusion_intermediate=64,
conf_hidden=256,
max_question_tokens=64,
max_option_tokens=48,
early_state_tokens=8,
)
model = GrootDecisionModel.tiny(cfg, vocab_size=1024, layers=2).eval()
model.save_pretrained(ckpt)
write_gbw(model, out / "model.gbw")
chunks, questions = build_request(rng, 1024)
blob = RequestBlob.build(chunks, questions)
RequestBlob.write(blob, out / "req.bin")
# ---- torch reference over the identical token structures
state = EncodedState(chunks=chunks, tokens=sum(len(c) - 2 for c in chunks))
eqs = questions
batch = collate([state], eqs, PAD, q_state=[0] * len(eqs))
with torch.no_grad():
res = model.forward(batch)
nopt = sum(q.n_options for q in questions)
nq = len(questions)
ref = np.zeros(nopt + nq + nopt, dtype="<f4")
o0 = 0
probs, conf, logits = res["probabilities"].numpy(), res["confidence_logit"].numpy(), res["logits"].numpy()
for qi, q in enumerate(questions):
k = q.n_options
ref[o0 : o0 + k] = probs[qi, :k]
ref[nopt + qi] = conf[qi]
ref[nopt + nq + o0 : nopt + nq + o0 + k] = logits[qi, :k]
o0 += k
np.save(out / "ref.npy", ref)
meta = {
"data_offset": data_offset(out / "model.gbw"),
"n_chunks": len(chunks),
"n_questions": nq,
"n_options": nopt,
"ref_len": int(ref.size),
}
(out / "meta.json").write_text(json.dumps(meta, indent=2) + "\n")
print(json.dumps(meta, indent=2))
if __name__ == "__main__":
main()
|