Buckets:
| """Our models on the Jevals suite (jevals.com, suite 0.1.0): the same 300 items per board, rebuilt from the source | |
| datasets and checked against the published state hashes, scored with Jevals' published rules. | |
| python jevals_eval.py build # CPU: rebuild and verify items -> data/jevals/items_*.jsonl | |
| python jevals_eval.py run --model checkpoints/... --name v2 # GPU: fast readout on all three boards | |
| python jevals_eval.py score # CPU: Decision Score, accuracy, ECE, bootstrap CIs, vs the board | |
| Boards: noul = PubMedQA (P(yes)), score = HelpSteer2 helpfulness (5 levels), choice = Banking77 (77 intents). | |
| Our readout is a System One answer: one forward pass, letter probabilities, no text. Prompts are our own typed format | |
| (the model's native interface); Jev is likewise called natively. Choice has 77 options and the readout has 26 | |
| letters, so Banking77 runs in two stages: the options (in a shuffled order) are split into chunks of <= 26 and read | |
| out per chunk; the top 8 of each chunk (24) are read out together, and those probabilities are the answer, with 0 on | |
| the other 53 (as Jevals' top-5 mode leaves unlisted options at the remainder, which is 0 when the listed sum to 1). | |
| Repeats follow Jevals: 5 per item, repeats 0 and 1 share an option order, 2-4 each get a new one (our seed; Jevals' | |
| order seed is not published). The readout is deterministic, so noul and score repeats are identical. | |
| Temperatures: per question type, fitted on our own calibration items for that model (never on Jevals items); | |
| raw (T = 1) scores are reported too. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import random | |
| import numpy as np | |
| from common import read_jsonl, softmax, write_jsonl | |
| D = "data/jevals" | |
| BOARDS = {"pubmedqa": "noul", "helpsteer2": "score", "banking77": "choice"} | |
| KEY_RENAME = {"banking77": {"text": "message"}, "pubmedqa": {"context.contexts": "context"}} # how Jevals names the state | |
| def js(x) -> str: | |
| return json.dumps(x, ensure_ascii=False, separators=(",", ":")) | |
| def build() -> None: | |
| from datasets import load_dataset | |
| for board, prim in BOARDS.items(): | |
| d = json.load(open(f"{D}/{board}.json")) | |
| ds = load_dataset(d["dataset"], d["config"], split=d["split"], revision=d["hf_revision"]) | |
| rows, bad = [], 0 | |
| for it in d["items"]: | |
| row = ds[int(it["row_idx"])] | |
| state = {} | |
| for fld in d["state_fields"]: | |
| a, _, b = fld.partition(".") | |
| state[KEY_RENAME.get(board, {}).get(fld, fld)] = row[a][b] if b else row[a] | |
| if hashlib.sha256(js(state).encode("utf-8")).hexdigest() != it["state_sha256"]: | |
| bad += 1 | |
| continue | |
| if prim == "noul": # our yes/no format: option 0 = true (yes), 1 = false (no) | |
| opts = [{"id": "true", "desc": d["criteria"]["true"]}, {"id": "false", "desc": d["criteria"]["false"]}] | |
| yes = d["options"][it["target"]] == "yes" | |
| gold, target = (0 if yes else 1), ([1.0, 0.0] if yes else [0.0, 1.0]) | |
| elif prim == "score": | |
| opts = [{"id": str(i), "desc": c} for i, c in enumerate(d["criteria"])] | |
| gold, target = it["target"], [float(i == it["target"]) for i in range(len(opts))] | |
| else: | |
| opts = [{"id": None, "desc": o} for o in d["options"]] | |
| gold, target = it["target"], [float(i == it["target"]) for i in range(len(opts))] | |
| rows.append({"id": it["item_id"], "set": f"jevals_{board}", "board": board, "qtype": prim, "evidence": js(state), | |
| "criterion": d["instructions"], "options": opts, "target": target, "gold": gold}) | |
| write_jsonl(f"{D}/items_{board}.jsonl", rows) | |
| print(json.dumps({"board": board, "items": len(rows), "hash_mismatch": bad}), flush=True) | |
| def run(args) -> None: | |
| from transformers import AutoTokenizer | |
| from vllm_utils import make_llm, readout | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| llm = make_llm(args.model, gpu_mem=args.gpu_mem) | |
| out = {} | |
| for board, prim in BOARDS.items(): | |
| items = read_jsonl(f"{D}/items_{board}.jsonl") | |
| if prim != "choice": | |
| lps = readout(llm, tok, items, [list(range(len(x["options"]))) for x in items]) | |
| out[board] = [{"id": x["id"], "repeats": [lp.tolist()] * 5} for x, lp in zip(items, lps)] | |
| continue | |
| rng = random.Random(20261031) | |
| orders = {x["id"]: [] for x in items} | |
| for x in items: # repeats 0-1 share an order; 2, 3, 4 each get a new one | |
| o0 = list(range(len(x["options"]))) | |
| rng.shuffle(o0) | |
| orders[x["id"]] = [o0, o0] + [rng.sample(o0, len(o0)) for _ in range(3)] | |
| res = {x["id"]: [None] * 5 for x in items} | |
| for r in (0, 2, 3, 4): | |
| # stage 1: chunks of <= 26 in the repeat's order | |
| sub_items, sub_orders, owner = [], [], [] | |
| for x in items: | |
| o = orders[x["id"]][r] | |
| n = -(-len(o) // 26) | |
| size = -(-len(o) // n) | |
| for c in range(n): | |
| chunk = o[c * size:(c + 1) * size] | |
| sub_items.append({**x, "options": [x["options"][j] for j in chunk]}) | |
| sub_orders.append(list(range(len(chunk)))) | |
| owner.append((x["id"], chunk)) | |
| lps = readout(llm, tok, sub_items, sub_orders) | |
| cand = {x["id"]: [] for x in items} | |
| for (iid, chunk), lp in zip(owner, lps): | |
| p = softmax(np.asarray(lp), 1.0) | |
| cand[iid] += [chunk[j] for j in np.argsort(-p)[:args.top]] | |
| # stage 2: the candidates together, in the order they came out of the repeat's shuffle | |
| s2_items, s2_cands = [], [] | |
| for x in items: | |
| o = orders[x["id"]][r] | |
| cs = sorted(cand[x["id"]], key=o.index) | |
| s2_items.append({**x, "options": [x["options"][j] for j in cs]}) | |
| s2_cands.append(cs) | |
| lps2 = readout(llm, tok, s2_items, [list(range(len(c))) for c in s2_cands]) | |
| for x, cs, lp in zip(items, s2_cands, lps2): | |
| res[x["id"]][r] = {"cands": cs, "lp": lp.tolist()} | |
| for x in items: | |
| res[x["id"]][1] = res[x["id"]][0] | |
| out[board] = [{"id": x["id"], "repeats": res[x["id"]]} for x in items] | |
| os.makedirs(f"{D}/runs", exist_ok=True) | |
| json.dump({"model": args.model, "name": args.name, "top": args.top, "boards": out}, open(f"{D}/runs/{args.name}.json", "w")) | |
| print(json.dumps({"saved": f"{D}/runs/{args.name}.json"}), flush=True) | |
| def think(args) -> None: | |
| """Thinking mode: one chain per item, then the readout after the chain (the deployed thinking path). Choice reasons | |
| over the 24 options the same model's fast readout shortlisted (repeat 0), so a gold outside the shortlist is wrong. | |
| An unfinished chain (no answer by the context limit) falls back to the fast readout, as the routed system would.""" | |
| from transformers import AutoTokenizer | |
| from common import split_chain | |
| from vllm_utils import make_llm, readout, think_sync | |
| fast = json.load(open(f"{D}/runs/{args.name}.json"))["boards"] | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| llm = make_llm(args.model, gpu_mem=args.gpu_mem) | |
| out = {} | |
| for board, prim in BOARDS.items(): | |
| items = read_jsonl(f"{D}/items_{board}.jsonl") | |
| fast_b = {r["id"]: r["repeats"][0] for r in fast[board]} | |
| cands = [fast_b[x["id"]]["cands"] if prim == "choice" else list(range(len(x["options"]))) for x in items] | |
| sub = [{**x, "options": [x["options"][j] for j in c]} for x, c in zip(items, cands)] | |
| orders = [list(range(len(c))) for c in cands] | |
| samples = [ss[0] for ss in think_sync(llm, tok, sub, orders, n=1, seed=17)] | |
| keep = [i for i, smp in enumerate(samples) if smp["closed"]] | |
| post = readout(llm, tok, [sub[i] for i in keep], [orders[i] for i in keep], | |
| think_chains=[split_chain(samples[i]["text"])[0] for i in keep]) | |
| post = dict(zip(keep, post)) | |
| rows = [] | |
| for i, (x, c, smp) in enumerate(zip(items, cands, samples)): | |
| if i in post: | |
| rep = {"cands": c, "lp": post[i].tolist()} if prim == "choice" else post[i].tolist() | |
| else: | |
| rep = fast_b[x["id"]] | |
| rows.append({"id": x["id"], "repeats": [rep] * 5, "tokens": smp["tokens"], "closed": smp["closed"], "answer": smp["answer"]}) | |
| out[board] = rows | |
| print(json.dumps({"board": board, "closed": len(keep), "of": len(items), | |
| "tokens_median": float(np.median([s["tokens"] for s in samples]))}), flush=True) | |
| json.dump({"model": args.model, "name": args.name + "_think", "mode": "think", "boards": out}, open(f"{D}/runs/{args.name}_think.json", "w")) | |
| print(json.dumps({"saved": f"{D}/runs/{args.name}_think.json"}), flush=True) | |
| def probs(rep, k: int, T: float) -> np.ndarray: | |
| if isinstance(rep, dict): # two-stage choice answer | |
| p = np.zeros(k) | |
| p[rep["cands"]] = softmax(np.asarray(rep["lp"]), T) | |
| return p | |
| return softmax(np.asarray(rep), T) | |
| def loss(p: np.ndarray, y: int, prim: str) -> float: | |
| if prim == "score": # ranked probability score over cumulative levels | |
| P, Y = np.cumsum(p)[:-1], (np.arange(len(p) - 1) >= y).astype(float) | |
| return float(np.sum((P - Y) ** 2) / (len(p) - 1)) | |
| e = np.zeros(len(p)); e[y] = 1.0 # multiclass Brier | |
| return float(np.sum((p - e) ** 2)) | |
| def ece10(conf, ok) -> float: | |
| conf, ok = np.asarray(conf), np.asarray(ok, dtype=float) | |
| b = np.minimum(9, np.floor(np.round(100 * conf) / 10)).astype(int) | |
| return float(sum((b == i).mean() * abs(ok[b == i].mean() - conf[b == i].mean()) for i in range(10) if (b == i).any())) | |
| def score_run(name: str, temps: dict | None) -> dict: | |
| R = json.load(open(f"{D}/runs/{name}.json")) | |
| out = {} | |
| for board, prim in BOARDS.items(): | |
| items = {x["id"]: x for x in read_jsonl(f"{D}/items_{board}.jsonl")} | |
| T = (temps or {}).get(prim, 1.0) | |
| rows = R["boards"][board] | |
| k = len(next(iter(items.values()))["options"]) | |
| prior = np.bincount([items[r["id"]]["gold"] for r in rows], minlength=k) / len(rows) | |
| L, Lp, conf, ok = [], [], [], [] | |
| for r in rows: | |
| y = items[r["id"]]["gold"] | |
| ps = [probs(rep, k, T) for rep in r["repeats"]] | |
| L.append(np.mean([loss(p, y, prim) for p in ps])) | |
| Lp.append(loss(prior, y, prim)) | |
| for p in ps: | |
| if prim == "noul": | |
| c = max(p[0], p[1]); pick = None if p[0] == p[1] else int(np.argmax(p)) | |
| else: | |
| c = float(p.max()); pick = int(np.argmax(p)) | |
| conf.append(c); ok.append(pick == y) | |
| L, Lp = np.array(L), np.array(Lp) | |
| ds = 100 * (1 - L.mean() / Lp.mean()) | |
| rng = np.random.default_rng(2000) | |
| bs = [] | |
| for _ in range(2000): | |
| i = rng.integers(0, len(L), len(L)) | |
| bs.append(100 * (1 - L[i].mean() / Lp[i].mean())) | |
| out[board] = {"primitive": prim, "ds": ds, "ci": [float(np.percentile(bs, 2.5)), float(np.percentile(bs, 97.5))], | |
| "accuracy": float(np.mean(ok)), "ece_10_bin": ece10(conf, ok), "T": T, "n": len(rows)} | |
| return out | |
| def temps_post_for_model(reason_file: str) -> dict: | |
| from common import fit_temperature | |
| cal = {x["id"]: x for x in read_jsonl("data/calib.jsonl")} | |
| rows = [r for r in read_jsonl(reason_file) if r["id"] in cal and r.get("post_lp") is not None] | |
| return {qt: fit_temperature([np.array(r["post_lp"]) for r in rows if cal[r["id"]]["qtype"] == qt], | |
| [cal[r["id"]]["gold"] for r in rows if cal[r["id"]]["qtype"] == qt]) for qt in ("choice", "noul", "score")} | |
| def temps_for_model(pred_file: str) -> dict: | |
| from common import fit_temperature | |
| cal = read_jsonl("data/calib.jsonl") | |
| lps = {r["id"]: np.array(r["lp"]) for r in read_jsonl(pred_file)} | |
| return {qt: fit_temperature([lps[x["id"]] for x in cal if x["qtype"] == qt], [x["gold"] for x in cal if x["qtype"] == qt]) | |
| for qt in ("choice", "noul", "score")} | |
| def score_all(args) -> None: | |
| board = json.load(open(f"{D}/board.json")) | |
| ours = {} | |
| for f in sorted(os.listdir(f"{D}/runs")): | |
| name = f[:-5] | |
| R = json.load(open(f"{D}/runs/{f}")) | |
| if name.endswith("_think"): | |
| pf = REASON.get(name[:-6]) | |
| temps = temps_post_for_model(pf) if pf and os.path.exists(pf) else None | |
| else: | |
| pf = PRED.get(name) | |
| temps = temps_for_model(pf) if pf and os.path.exists(pf) else None | |
| ours[name] = {"model": R["model"], "temps": temps, "raw": score_run(name, None), | |
| "calibrated": score_run(name, temps) if temps else None} | |
| rows = [] | |
| for r in board["rows"]: | |
| rows.append({"board": {v: k for k, v in BOARDS.items()}[r["primitive"]], "name": r["display_name"], "ds": r["decision_score"], | |
| "ci": [r["ci_low"], r["ci_high"]], "accuracy": r["accuracy"], "ece": r["ece_10_bin"], "p95_ms": r["p95_ms"], | |
| "cost_per_1k": r["cost_per_1k"], "interface": r["interface"]}) | |
| json.dump({"ours": ours, "board": rows, "release": board["release_id"], "suite": board["suite_version"]}, | |
| open("results/jevals.json", "w"), indent=1) | |
| for b in BOARDS: | |
| print(f"\n== {b} ({BOARDS[b]})") | |
| tab = [(x["name"], x["ds"], x["ci"], x["accuracy"], x["ece"], "board") for x in rows if x["board"] == b] | |
| for n, o in ours.items(): | |
| s = (o["calibrated"] or o["raw"])[b] | |
| tab.append((f"OURS {n}" + (" (cal)" if o["calibrated"] else " (raw)"), s["ds"], s["ci"], s["accuracy"], s["ece_10_bin"], "ours")) | |
| if o["calibrated"]: | |
| s = o["raw"][b]; tab.append((f"OURS {n} (raw T=1)", s["ds"], s["ci"], s["accuracy"], s["ece_10_bin"], "ours")) | |
| for n, ds, ci, acc, ece, _ in sorted(tab, key=lambda t: -t[1]): | |
| print(f" {n:34s} DS {ds:6.1f} [{ci[0]:6.1f},{ci[1]:6.1f}] acc {acc:.3f} ECE {ece:.3f}") | |
| # our readout prediction files (with calibration items) per run name, for the temperatures | |
| PRED = {"base": "results/predictions/readout_base.jsonl", "v1_joint": "results/predictions/readout_joint.jsonl", | |
| "v2_joint": "runs/v2/results/predictions/readout_v2_joint.jsonl", "rl60": "runs/v2/results/predictions/readout_rl60.jsonl", | |
| "rl40": "runs/v2/results/predictions/readout_rl40.jsonl", "base_2b": "runs/v2_2b/results/predictions/readout_base2b.jsonl", | |
| "v2_2b": "runs/v2_2b/results/predictions/readout_v2b.jsonl", "v2_2b_ro": "runs/v2_2b/results/predictions/readout_v2b_ro.jsonl"} | |
| # our thinking-mode prediction files (with calibration chains and post-chain readouts), for post-chain temperatures | |
| REASON = {"base": "results/predictions/reason_base.jsonl", "v1_joint": "results/predictions/reason_joint.jsonl", | |
| "v2_joint": "runs/v2/results/predictions/reason_v2_joint.jsonl", "rl60": "runs/v2/results/predictions/reason_rl60.jsonl", | |
| "rl40": "runs/v2/results/predictions/reason_rl40.jsonl", "v2_2b": "runs/v2_2b/results/predictions/reason_v2b.jsonl", | |
| "base_2b": "runs/v2_2b/results/predictions/reason_base2b.jsonl", "v2_2b_ro": "runs/v2_2b/results/predictions/reason_v2b_ro.jsonl"} | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("cmd", choices=["build", "run", "think", "score"]) | |
| ap.add_argument("--model") | |
| ap.add_argument("--name") | |
| ap.add_argument("--top", type=int, default=8, help="choice: candidates kept per stage-1 chunk") | |
| ap.add_argument("--gpu-mem", type=float, default=0.92, help="0.92 of the card: ~6 GB more KV cache than 0.85, nothing else on the GPU") | |
| ap.add_argument("--max-seqs", type=int, default=384, help="think: concurrent chains (256 before; more KV now)") | |
| ap.add_argument("--batched-tokens", type=int, default=32768, help="prefill tokens per engine step (readouts are prefill-bound)") | |
| args = ap.parse_args() | |
| # scheduling knobs only (vllm_utils reads them when the engine starts); they never change what the model computes | |
| os.environ.setdefault("MAX_BATCHED_TOKENS", str(args.batched_tokens)) | |
| if args.cmd == "think": | |
| os.environ.setdefault("MAX_NUM_SEQS", str(args.max_seqs)) | |
| {"build": lambda: build(), "run": lambda: run(args), "think": lambda: think(args), "score": lambda: score_all(args)}[args.cmd]() | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 16.9 kB
- Xet hash:
- 7ba644874f466fa9bf0215ea17ba38582b84f59b69aff98200d7bafa65ef9ff3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.