Buckets:
| """Evaluate a model on Core (or dev) at several K from one trajectory per item. | |
| The head is applied after every pass 1..Kmax, so per-pass probabilities are saved (E4 adaptive stopping is | |
| computed offline from them). Weights run in bf16 like S1's published eval. Default max length covers every | |
| Core item (no item is skipped); `in_s1_protocol` marks items within S1's 4,096-token protocol. | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from collections import defaultdict | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from common import CORE14_EXCLUDE, DATA, GROUPS, PROC, S1_PATH, Renderer, load_canonical, load_core, tokenizer_and_mapping # noqa | |
| from model import ThinkDecide, collate, load_backbone # noqa: E402 | |
| from bench_local import metrics # noqa: E402 (jev-dllm's own metric definitions) | |
| def build_model(path, n_think, use_head, device, mapping, dtype=torch.bfloat16, fb_mode=None): | |
| base = load_backbone(S1_PATH, dtype, device) | |
| if fb_mode is None: # read the arm's own config when available | |
| cfg = os.path.join(os.path.dirname(path or ""), "config.json") | |
| fb_mode = json.load(open(cfg)).get("fb_mode", "hidden") if path and os.path.exists(cfg) else "hidden" | |
| model = ThinkDecide(base, mapping["Yes"], mapping["No"], n_think=n_think, use_head=use_head, | |
| fb_mode=fb_mode).to(device, dtype) | |
| if path and path != "s1": | |
| from safetensors.torch import load_file | |
| sd = load_file(path, device=str(device)) | |
| missing, unexpected = model.load_state_dict(sd, strict=False) | |
| missing = [m for m in missing if m != "base.lm_head.weight" and not m.startswith("bb.")] | |
| assert not missing and not unexpected, (missing, unexpected) | |
| model.base.tie_weights() | |
| assert model.base.lm_head.weight.data_ptr() == model.bb.embed_tokens.weight.data_ptr() | |
| return model.eval() | |
| def batches(items, budget, max_items): | |
| order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"])) | |
| out, cur = [], [] | |
| for i in order: | |
| L = len(items[i]["ids"]) | |
| if cur and (L * (len(cur) + 1) > budget or len(cur) >= max_items): | |
| out.append(cur) | |
| cur = [] | |
| cur.append(i) | |
| if cur: | |
| out.append(cur) | |
| return out | |
| def probs_from_scores(s_row, prim, n): | |
| s = s_row[:n].astype(np.float64) | |
| if prim == "noul": | |
| p1 = 1.0 / (1.0 + np.exp(-s[0])) | |
| return np.array([1.0 - p1, p1]) | |
| e = np.exp(s - s.max()) | |
| return e / e.sum() | |
| def run(model, items, kmax, device, pad_id, budget, max_items, read_all=True, batch_ids=None): | |
| reads = list(range(1, kmax + 1)) if read_all else [kmax] | |
| per_pass = [None] * len(items) | |
| t0 = time.time() | |
| if batch_ids: # explicit batch composition (parity check against jev-dllm's batches.jsonl) | |
| where = {x["id"]: i for i, x in enumerate(items)} | |
| plan = [[where[i] for i in ids if i in where] for ids in batch_ids] | |
| plan = [b for b in plan if b] | |
| else: | |
| plan = batches(items, budget, max_items) | |
| for bi, idx in enumerate(plan): | |
| mb = [items[i] for i in idx] | |
| batch = collate(mb, device, pad_id) | |
| outs = model(batch, K=kmax, read_at=reads) | |
| outs = [o.float().cpu().numpy() for o in outs] | |
| for j, i in enumerate(idx): | |
| x = items[i] | |
| n = 1 if x["primitive"] == "noul" else len(x["pos"]) | |
| per_pass[i] = [probs_from_scores(o[j], x["primitive"], n if x["primitive"] != "noul" else 1) | |
| for o in outs] | |
| torch.cuda.synchronize() | |
| return per_pass, reads, time.time() - t0 | |
| def summarize(rows, key): | |
| """rows have gold_idx and probs[key]; jev-dllm metrics (acc, ECE15, Brier, NLL) per suite + groups.""" | |
| by_suite = defaultdict(list) | |
| for r in rows: | |
| by_suite[r["suite"]].append((r["gold_idx"], r["probs"][key])) | |
| suites = {s: metrics(v) for s, v in by_suite.items()} | |
| groups = {} | |
| for g, names in GROUPS.items(): | |
| names = [n for n in names if n in suites] | |
| if names: | |
| groups[g] = {m: float(np.mean([suites[n][m] for n in names])) for m in ("accuracy", "ece", "nll", "brier")} | |
| core14 = [s for s in suites if s not in CORE14_EXCLUDE] | |
| groups["core17_macro"] = {m: float(np.mean([suites[s][m] for s in suites])) for m in ("accuracy", "ece")} | |
| groups["core14_macro"] = {m: float(np.mean([suites[s][m] for s in core14])) for m in ("accuracy", "ece")} | |
| groups["reasoning_no_mnli"] = {m: float(np.mean([suites[s][m] for s in GROUPS["reasoning"] | |
| if s != "nimble_multinli" and s in suites])) | |
| for m in ("accuracy", "ece")} | |
| return dict(suites=suites, groups=groups) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", default="s1", help="model.safetensors path, or s1") | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--n-think", type=int, default=0) | |
| ap.add_argument("--no-head", action="store_true") | |
| ap.add_argument("--kmax", type=int, default=1) | |
| ap.add_argument("--split", choices=["core", "dev", "proc_test", "proc_dev"], default="core") | |
| ap.add_argument("--compile", action="store_true") | |
| ap.add_argument("--max-length", type=int, default=16384) | |
| ap.add_argument("--budget", type=int, default=65536, help="padded tokens per forward") | |
| ap.add_argument("--max-items", type=int, default=32) | |
| ap.add_argument("--limit-per-suite", type=int, default=0) | |
| ap.add_argument("--batches-from", default="", help="jev-dllm batches.jsonl: reuse its batch composition") | |
| a = ap.parse_args() | |
| os.makedirs(a.out, exist_ok=True) | |
| device = torch.device("cuda") | |
| tok, mapping, mask_id, mask_text = tokenizer_and_mapping() | |
| rend = Renderer(tok, mask_id, mask_text, a.n_think) | |
| rend0 = Renderer(tok, mask_id, mask_text, 0) | |
| if a.split == "core": | |
| items, unsupported = load_core(rend, a.max_length, limit_per_suite=a.limit_per_suite) | |
| lens0 = {x["id"]: len(x["ids"]) for x in load_core(rend0, 10**9, limit_per_suite=a.limit_per_suite)[0]} | |
| else: | |
| path = {"dev": DATA + "/dev.jsonl", "proc_test": PROC + "/test.jsonl", "proc_dev": PROC + "/dev.jsonl"}[a.split] | |
| items, _ = load_canonical(path, rend, 10**9) | |
| unsupported = [] | |
| for i, x in enumerate(items): | |
| x.update(id="%s|%s" % (x["case_id"], x["view_id"]), | |
| suite=("dev:" if a.split == "dev" else "") + x["source"], | |
| gold_idx=int(np.argmax(x["target"])) if x["kind"] == "hard" else None) | |
| lens0 = {} | |
| model = build_model(a.model, a.n_think, not a.no_head, device, mapping) | |
| if a.compile: | |
| model.compile_layers() | |
| batch_ids = [json.loads(l)["ids"] for l in open(a.batches_from)] if a.batches_from else None | |
| per_pass, reads, secs = run(model, items, a.kmax, device, tok.pad_token_id, a.budget, a.max_items, | |
| batch_ids=batch_ids) | |
| rows = [] | |
| for x, pp in zip(items, per_pass): | |
| r = {k: x[k] for k in ("id", "suite", "primitive", "gold_idx") if k in x} | |
| r["length"] = len(x["ids"]) | |
| if lens0: | |
| r["in_s1_protocol"] = lens0[x["id"]] <= 4096 | |
| if a.split != "core": | |
| r["target"] = x["target"] | |
| r["kind"] = x["kind"] | |
| r["level"] = x.get("level") | |
| r["probs"] = {str(k): [round(float(v), 7) for v in p] for k, p in zip(reads, pp)} | |
| rows.append(r) | |
| with open(os.path.join(a.out, "predictions.jsonl"), "w") as f: | |
| for r in rows: | |
| f.write(json.dumps(r) + "\n") | |
| hard_rows = [r for r in rows if r.get("gold_idx") is not None] | |
| report = dict(model=a.model, n_think=a.n_think, head=not a.no_head, kmax=a.kmax, split=a.split, | |
| n_items=len(rows), unsupported=len(unsupported), forward_seconds=round(secs, 1), | |
| by_K={str(k): summarize(hard_rows, str(k)) for k in reads}) | |
| if a.split.startswith("proc"): | |
| def acc(rows, k): | |
| return float(np.mean([np.argmax(r["probs"][k]) == r["gold_idx"] for r in rows])) if rows else None | |
| report["proc"] = {str(k): dict( | |
| all=acc(hard_rows, str(k)), | |
| by_level={str(l): acc([r for r in hard_rows if r["level"] == l], str(k)) for l in range(5)}, | |
| by_task={t: acc([r for r in hard_rows if r["suite"] == t], str(k)) for t in sorted({r["suite"] for r in hard_rows})}) | |
| for k in reads} | |
| for k in reads: | |
| pr = report["proc"][str(k)] | |
| print("K=%d proc all %.2f | " % (k, 100 * pr["all"]) + " ".join("L%s %.2f" % (l, 100 * v) for l, v in pr["by_level"].items())) | |
| if a.split == "core": | |
| sub = [r for r in hard_rows if r["in_s1_protocol"]] | |
| report["s1_protocol_subset"] = dict(n=len(sub), by_K={str(k): summarize(sub, str(k)) for k in reads}) | |
| json.dump(report, open(os.path.join(a.out, "report.json"), "w"), indent=1) | |
| print(json.dumps(dict(event="eval_done", out=a.out, n=len(rows), unsupported=len(unsupported), | |
| seconds=round(secs, 1)))) | |
| for k in reads: | |
| g = report["by_K"][str(k)]["groups"] | |
| print("K=%d " % k + " ".join("%s %.2f" % (n, 100 * v["accuracy"]) for n, v in g.items())) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 9.39 kB
- Xet hash:
- 50c676eaa7fa6949b38edc7ab55601639e538bbd1075c924b5860486930ab971
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.