#!/usr/bin/env python3 """Arm 11 -- train the estate's own decision model on the estate's own labels. python3 run_arm11.py --repo ~/bench-ownjev --out rows-addenda Pre-registered in README "Addenda arms 4-11". Trains the Apache-2.0 openjev framework's ModernBERT-base ENCODER variant on three estate sets, then scores the held-out 30 % against three comparators on the SAME held-out items. WHAT THIS FILE REFUSES TO DO, all of it pre-registered. * **It never trains on a held-out row.** The split was written to disk by `build_ownjev_tasks.py` before any training step; this file re-reads `heldout_ids.json` and ASSERTS that the train and dev files share no uid with the test file before it will start a run. * **No number for a held-out set under 12 items.** The cell says "too small". * **A comparator that does not cover the held-out items is left empty, with the reason** -- never back-filled from the full set, and never replaced by a figure from a neighbouring task. * It touches inferencebox for nothing, and it runs only after the vLLM server has stopped, so the cards are its own and the training seconds are real. THE COMPARATORS ARE FILTERED, NOT RE-RUN. OpenJev-FP8 and gemma 4 already answered every one of these items zero-shot in arms 2/4/5 and arm 10; their rows are on disk. This file selects the held-out subset of those rows by the same uid the split used, so the three numbers in a row of the table are three models answering the SAME thirty-odd items. """ from __future__ import annotations import argparse import json import os import subprocess import sys import time HERE = os.path.dirname(os.path.abspath(__file__)) MIN_HELDOUT = 12 #: task dir -> (the estate rows that answered the same items zero-shot) #: Each entry: (rows dir, arm name, row file task letter, how to read a row's #: uid, how to read a row's predicted label in THIS task's label space). COMPARATORS = { "longtable_voices": { "question": "speaker", "openjev": ("rows", "openjev-fp8-readout", "c"), "gemma4": ("rows-addenda", "gemma4-fp8-readout", "c"), "map": lambda row: row.get("choice"), }, "doorman_gate": { "question": "refuse", "openjev": ("rows-addenda", "openjev-fp8-readout", "d"), "gemma4": ("rows-addenda", "gemma4-fp8-readout", "d"), # the doorman arms answer admit/refuse; the trained task is a noul on # "should the table REFUSE this", so refuse -> true "map": lambda row: "true" if row.get("choice") == "refuse" else "false", }, "fieldexam_grounded": { "question": "answerable", "openjev": ("rows-addenda", "openjev-fp8-readout", "f"), "gemma4": ("rows-addenda", "gemma4-fp8-readout", "f"), "map": lambda row: "true" if row.get("choice") == "yes" else "false", }, } def uid_of_rows(rows): """Re-derive the split's uid from a row file, the same way the split did.""" import collections seen = collections.Counter() out = [] for r in rows: seen[r["id"]] += 1 n = seen[r["id"]] out.append(r["id"] if n == 1 else "%s#%d" % (r["id"], n)) return out def read_rows(rel, arm, task): p = os.path.join(HERE, rel, "%s.%s.jsonl" % (arm, task)) if not os.path.exists(p): return None return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()] def brier(prob_dicts, golds): tot = 0.0 for p, g in zip(prob_dicts, golds): tot += sum((v - (1.0 if k == g else 0.0)) ** 2 for k, v in p.items()) return tot / len(golds) if golds else None def dir_size_mb(path): n = 0 for root, _d, files in os.walk(path): for f in files: n += os.path.getsize(os.path.join(root, f)) return n / 1e6 #: benchbox has no C compiler and no CUDA toolkit, and a lane has no sudo to #: install either. Arm 2 already solved exactly this for vLLM's triton JIT with #: a shim on PATH (`serve.sh`), and triton needs the same compiler here to build #: its kernels for the encoder. Reusing arm 2's shim rather than inventing a #: second one keeps one answer to one box problem. CUDA_VISIBLE_DEVICES pins #: card 0: the framework's device selection is a bare `cuda`, which is `cuda:0`, #: and a bench that did not say which board it trained on would be repeating the #: trap D-20260921-001 exists to prevent. SHIM = "/workshop/bench-arm2/shim" def train_env(): return dict(os.environ, HF_HOME="/workshop/hf-cache", CUDA_VISIBLE_DEVICES="0", PATH=SHIM + ":" + os.environ.get("PATH", ""), CC=os.path.join(SHIM, "cc"), CXX=os.path.join(SHIM, "c++")) def train_one(repo, task, epochs, bs, init_from=None): env = train_env() cmd = [os.path.join(repo, ".venv/bin/python"), "scripts/train.py", "--task", "tasks/%s" % task, "--epochs", str(epochs), "--bs", str(bs)] if init_from: # reading B: the framework's own central claim -- the same data and the # same time, started from its published general encoder instead of the # raw backbone. `--overwrite` because reading A already wrote a # checkpoint of this name; the two readings are kept apart by writing # their reports and held-out rows under different names, not by sharing # a checkpoint path. cmd += ["--init-from", init_from, "--overwrite"] print("ANNOUNCE box benchbox · card 0 (one RTX 3090, 250 W) · arm 11 · " "train encoder `%s` · %s · %s" % (task, " ".join(cmd[1:]), time_now()), flush=True) t0 = time.monotonic() p = subprocess.run(cmd, cwd=repo, env=env, capture_output=True, text=True) secs = time.monotonic() - t0 return secs, p def time_now(): import datetime as dt return dt.datetime.now(dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") def evaluate(repo, task, qid): """Score the trained checkpoint on the held-out test.jsonl.""" code = r''' import json, sys, torch sys.path.insert(0, %r) from openjev import DecisionModel, Task task = %r; qid = %r t = Task.load("tasks/" + task, "test") m = DecisionModel.from_pretrained("checkpoints/%%s/model.pt" %% task) states = [e.state for e in t.examples] outs = m.answer_batch(states, t.questions) rows = [] for e, o in zip(t.examples, outs): a = o[qid] gold = e.answers[qid] gold_s = ("true" if gold is True else "false" if gold is False else str(gold)) rows.append({"uid": e.meta.get("uid"), "id": e.meta.get("id"), "gold": gold_s, "pred": a.label, "probs": {k: float(v) for k, v in a.probabilities.items()}}) print("@@@" + json.dumps(rows)) ''' % (repo, task, qid) env = train_env() p = subprocess.run([os.path.join(repo, ".venv/bin/python"), "-c", code], cwd=repo, env=env, capture_output=True, text=True) for line in p.stdout.splitlines(): if line.startswith("@@@"): return json.loads(line[3:]), p return None, p def main(argv=None) -> int: ap = argparse.ArgumentParser() ap.add_argument("--repo", default=os.path.expanduser("~/bench-ownjev")) ap.add_argument("--out", default=os.path.join(HERE, "rows-addenda")) ap.add_argument("--epochs", type=int, default=6) ap.add_argument("--bs", type=int, default=8) ap.add_argument("--task", action="append") ap.add_argument("--init-from", default=None, help="reading B: warm-start from a published checkpoint") ap.add_argument("--tag", default="arm11", help="prefix for this reading's output files") a = ap.parse_args(argv) os.makedirs(a.out, exist_ok=True) tasks = a.task or list(COMPARATORS) reports = [] for task in tasks: spec = COMPARATORS[task] tdir = os.path.join(a.repo, "tasks", task) split = json.load(open(os.path.join(tdir, "heldout_ids.json"), encoding="utf-8")) test_ids = set(split["test_ids"]) # THE LEAK CHECK, before a single training step overlap = test_ids & (set(split["train_ids"]) | set(split["dev_ids"])) if overlap: raise SystemExit("arm 11: %s leaks %d held-out uids into training: %s" % (task, len(overlap), sorted(overlap)[:5])) rep = {"task": task, "question": spec["question"], "started": time_now(), "reading": ("B: warm-started from %s" % a.init_from) if a.init_from else "A: from the raw ModernBERT-base backbone", "heldout_n": len(test_ids), "split": split["counts"], "seed": split["seed"], "leak_check": "passed: train/dev share no uid with test"} if len(test_ids) < MIN_HELDOUT: rep["result"] = "too small" rep["reason"] = ("held-out set is %d items, under the pre-registered " "floor of %d; no rate, no Brier, no comparator delta" % (len(test_ids), MIN_HELDOUT)) print("%s: TOO SMALL (%d held out)" % (task, len(test_ids)), flush=True) reports.append(rep) continue secs, proc = train_one(a.repo, task, a.epochs, a.bs, a.init_from) rep["train_seconds"] = secs rep["train_rc"] = proc.returncode rep["train_tail"] = proc.stdout[-1500:] if proc.stdout else proc.stderr[-1500:] if proc.returncode != 0: rep["result"] = "training failed" rep["error"] = (proc.stderr or "")[-2000:] print("%s: TRAINING FAILED rc=%d" % (task, proc.returncode), flush=True) reports.append(rep) continue ckpt = os.path.join(a.repo, "checkpoints", task, "model.pt") rep["model_bytes"] = os.path.getsize(ckpt) if os.path.exists(ckpt) else None rep["model_mb"] = (rep["model_bytes"] / 1e6) if rep["model_bytes"] else None rows, ep = evaluate(a.repo, task, spec["question"]) if rows is None: rep["result"] = "eval failed" rep["error"] = (ep.stderr or "")[-2000:] reports.append(rep) continue n = len(rows) corr = sum(1 for r in rows if r["pred"] == r["gold"]) rep.update(trained={"n": n, "correct": corr, "accuracy": corr / n, "brier": brier([r["probs"] for r in rows], [r["gold"] for r in rows])}) with open(os.path.join(a.out, "%s.%s.heldout.jsonl" % (a.tag, task)), "w", encoding="utf-8") as fh: for r in rows: fh.write(json.dumps(r, ensure_ascii=False) + "\n") gold_by_uid = {r["uid"]: r["gold"] for r in rows} # the floor: the majority class ON THE HELD-OUT SET import collections cnt = collections.Counter(gold_by_uid.values()) rep["floor"] = {"kind": "majority class on the held-out set", "value": max(cnt.values()) / n, "counts": dict(cnt)} # the two zero-shot comparators, filtered to the held-out uids for who in ("openjev", "gemma4"): rel, arm, tl = spec[who] src = read_rows(rel, arm, tl) if not src: rep[who] = {"result": "no rows", "reason": "%s/%s.%s.jsonl is not on disk, so this " "comparator is empty rather than filled " "from another set" % (rel, arm, tl)} continue uids = uid_of_rows(src) sel = [(u, r) for u, r in zip(uids, src) if u in gold_by_uid] missing = sorted(set(gold_by_uid) - set(u for u, _ in sel)) if missing: rep[who] = {"result": "does not cover the held-out set", "covered": len(sel), "missing": len(missing), "reason": "left empty rather than back-filled"} continue c = sum(1 for u, r in sel if spec["map"](r) == gold_by_uid[u]) rep[who] = {"n": len(sel), "correct": c, "accuracy": c / len(sel), "source": "%s/%s.%s.jsonl, filtered to the held-out uids" % (rel, arm, tl)} rep["ended"] = time_now() print("%-20s heldout n=%d trained %.1f%% openjev %s gemma4 %s floor %.1f%% " "%.0fs %.0f MB" % (task, n, 100 * rep["trained"]["accuracy"], ("%.1f%%" % (100 * rep["openjev"]["accuracy"])) if "accuracy" in rep.get("openjev", {}) else "—", ("%.1f%%" % (100 * rep["gemma4"]["accuracy"])) if "accuracy" in rep.get("gemma4", {}) else "—", 100 * rep["floor"]["value"], secs, rep["model_mb"] or 0), flush=True) reports.append(rep) with open(os.path.join(a.out, "%s.summary.json" % a.tag), "w", encoding="utf-8") as fh: json.dump(reports, fh, indent=1, sort_keys=True, default=str) return 0 if __name__ == "__main__": raise SystemExit(main())