File size: 1,467 Bytes
7bf8323 | 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 | import sys, torch
from transformers import AutoModel, AutoTokenizer
from cases import CASES
import pack, q64ctx
from eos_graph import load_chunked, eos_inputs, with_slots, chunk_inputs
S, P = int(sys.argv[1]), int(sys.argv[2])
ref = AutoModel.from_pretrained("eos", trust_remote_code=True, device="mps")
tok = AutoTokenizer.from_pretrained("eos"); pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
g = load_chunked("eos", S, P); worst = 0
cases = CASES + [(q64ctx.state, dict(list(q64ctx.Q.items())[:9]))]
for state, qs in cases:
r = ref.system_one(state=state, questions=qs)["answers"]
jobs = pack.rows(tok, state, qs)
x, cand, qry, owner, used = eos_inputs(g.cfg, [e for _, _, e in jobs], S, P, pad)
x = with_slots(x, cand, qry, 64)
x["seg_chunks"], x["cont"], x["last_seg"] = chunk_inputs(x["segment"])
names = ("input_ids", "cos", "sin", "valid", "tail_onehot", "segment", "lag_keep", "lag_tail", "seg_chunks", "cont", "last_seg", "cand_idx", "query_idx")
with torch.no_grad():
lg = g(*[torch.from_numpy(x[k]) for k in names])
a = pack.answers(jobs, lg.tolist(), owner, None)
d = 0
for q in qs:
pr = r[q].get("probabilities", {"true": r[q].get("noul")}); pa = a[q].get("probabilities", {"true": a[q].get("noul")})
d = max(d, max(abs(pr[k] - pa[k]) for k in pr))
worst = max(worst, d); print(f"{len(qs)} questions, packed {used}, max dp {d:.2e}")
print("WORST", worst)
|