| 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) |
|
|
