Buckets:
| """cwm3/aide/eval/kv2_gate.py — mechanical decisions of the K-v2 programme (docs/decisions.md 2026-09-14 row "K-v2 programme — S0 | |
| pre-registration", aide/docs/protocol.md C23). Every decision is computed from pred_val.npz of finished runs + `harness.load_split("val")` | |
| (+ `metrics.hm_mask` on the val future block) and written as UTC-stamped json in aide/results/kv2/. Nothing is typed by hand; the val HM | |
| subset is never a ranking statistic; test numbers are never read here. | |
| python aide/eval/kv2_gate.py s1 # S1 val screen (seed 42): 9 K-v2 configs vs cwm3d_kp_s42_v2 -> s1_decision.json | |
| python aide/eval/kv2_gate.py s1b <cfg> # S1b (seed 123) for one survivor -> s1b_<cfg>_decision.json | |
| python aide/eval/kv2_gate.py s3 # S3 (i) MB grid argmin -> s3_decision.json | |
| python aide/eval/kv2_gate.py pretrain-check [<cfg>] # dry trunk-only load of K's pretraining checkpoint into the K-v2 model (CPU) | |
| python aide/eval/kv2_gate.py flags kv2|mb <cfg> # trainer flags of a config (used by aide/train/run_kv2_program.sh) | |
| python aide/eval/kv2_gate.py name kv2|mb <cfg> <seed> [pre] | name kv2pre <cfg> (per-architecture pretraining run name) | |
| python aide/eval/kv2_gate.py read <json-stem> <key> # e.g. read s1 survivors -> "full nospatial" | |
| Exit codes: 0 decision complete; 3 decision written but incomplete (pending runs); 4 model / trainer not ready (pretrain-check); 2 usage. | |
| Rules (verbatim from the pre-registration): | |
| S1 statistic = paired per-clip Δ val All (config − K); SURVIVE iff mean Δ ≤ −0.30 px; ≤ 2 survivors ranked by mean Δ, candidates within | |
| 0.10 px are ties -> fewer components wins (components = enabled optional modules of {spatial, lean head, hazard BCE, hazard gating, | |
| physics+NLL}); HM guard (not a selection statistic): KILL iff Δ val HM (40 clips) mean > +3.0 AND CI excl. 0. | |
| S1b CONTINUE iff Δ val(seed 123) < 0 AND mean(Δ val 42, Δ val 123) ≤ −0.25. | |
| S3 MB-best = argmin val All over the 10 grid points; ties within 0.10 px -> fewer parameters. | |
| pretrain-check: reuse K's checkpoint iff n_skipped_shape == 0 AND every TRUNK_PREFIXES tensor of the K-v2 model was loaded. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import argparse # noqa: E402 | |
| import glob # noqa: E402 | |
| import json # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.eval import metrics as M # noqa: E402 | |
| from aide.eval.harness import align, load_split # noqa: E402 | |
| RUNS = P.RESULTS / "runs" | |
| KV2 = P.RESULTS / "kv2" | |
| K_CKPT = RUNS / "cwm3d_kp_batonpre_pretrain_s123" / "ckpt_best.pt" | |
| MB_CKPT = RUNS / "motionbert_kp_batonpre_pretrain_s123" / "ckpt_best.pt" | |
| SEEDS5 = (42, 123, 456, 789, 1011) | |
| N_BOOT, SEED = 10000, 0 | |
| # thresholds (pre-registered) | |
| S1_THRESH, S1_TIE, S1_MAX_SURVIVORS, HM_GUARD = -0.30, 0.10, 2, 3.0 | |
| S1B_THRESH = -0.25 | |
| S3_TIE = 0.10 | |
| RULE_S1 = (f"paired per-clip Δ val All (config − cwm3d_kp_s42_v2), n = 405, {N_BOOT} clip-bootstrap seed {SEED}; SURVIVE iff mean Δ ≤ {S1_THRESH} px; " | |
| f"≤ {S1_MAX_SURVIVORS} survivors ranked by mean Δ, ties within {S1_TIE} px -> fewer components (full/q060/q090/lam03 = 5, nospatial/khead/hazaux/plain = 4, nohaz = 3); " | |
| f"HM guard (never a ranking statistic): KILL iff Δ val HM (40 clips) mean > +{HM_GUARD} AND CI excl. 0") | |
| RULE_S1B = f"CONTINUE to seeds 456/789/1011 iff Δ val(seed 123) < 0 AND mean(Δ val 42, Δ val 123) ≤ {S1B_THRESH} px; else 'val-killed at seed 123'" | |
| RULE_S3 = f"MB-best = argmin val All over the 10 grid points (9 new + existing default); ties within {S3_TIE} px -> fewer parameters" | |
| RULE_PRE = "reuse cwm3d_kp_batonpre_pretrain_s123/ckpt_best.pt iff dry train_kp.load_init_ckpt(parts='trunk') reports n_skipped_shape == 0 AND zero missing TRUNK_PREFIXES tensors; else own pretraining cwm3d_v2_batonpre_pretrain_s123" | |
| # ----------------------------------------------------------------------------- configurations (single source of truth for the launcher, the gate and the report) | |
| # cfg -> (trainer flags, number of enabled optional components) | |
| KV2_CONFIGS = { | |
| "full": ([], 5), | |
| "nospatial": (["--no-v2-spatial"], 4), | |
| "khead": (["--no-lean-head"], 4), | |
| "nohaz": (["--no-hazard"], 3), | |
| "hazaux": (["--haz-mod", "0"], 4), | |
| "plain": (["--no-physics", "--w-nll", "0"], 4), | |
| "q060": (["--haz-quantile", "0.60"], 5), | |
| "q090": (["--haz-quantile", "0.90"], 5), | |
| "lam03": (["--haz-weight", "0.3"], 5), | |
| } | |
| KV2_MODEL = "cwm3d_v2" | |
| # cfg -> (model name, extra trainer flags); "default" = the existing motionbert_kp_s<k>_v2 runs | |
| MB_GRID = { | |
| "default": ("motionbert_kp", []), | |
| "d192": ("motionbert_kp_d192", []), | |
| "L4": ("motionbert_kp_L4", []), | |
| "lr1e-4": ("motionbert_kp", ["--lr", "1e-4"]), | |
| "lr1e-3": ("motionbert_kp", ["--lr", "1e-3"]), | |
| "wd0.01": ("motionbert_kp", ["--wd", "0.01"]), | |
| "wd0.1": ("motionbert_kp", ["--wd", "0.1"]), | |
| "do0.1": ("motionbert_kp", ["--dropout", "0.1"]), | |
| "d192_L4": ("motionbert_kp_d192_L4", []), | |
| "d192_L4_do0.1": ("motionbert_kp_d192_L4", ["--dropout", "0.1"]), | |
| } | |
| MB_NEW = tuple(c for c in MB_GRID if c != "default") | |
| def kv2_run(cfg: str, seed: int, pre: bool = False) -> str: | |
| tag = "" if cfg == "full" else f"_{cfg}" | |
| return f"cwm3d_v2{tag}_batonpre_trunkp123_s{seed}_v2" if pre else f"cwm3d_v2{tag}_s{seed}_v2" | |
| def mb_run(cfg: str, seed: int, pre: bool = False) -> str: | |
| tag = "" if cfg == "default" else f"_{cfg}" | |
| return f"motionbert_kp{tag}_batonpre_trunkp123_s{seed}_v2" if pre else f"motionbert_kp{tag}_s{seed}_v2" | |
| def kv2_pretrain_run(cfg: str) -> str: | |
| """Trunk-only pretraining run per K-v2 trunk ARCHITECTURE (coordinator 2026-09-14, after S1): `nospatial` has no spatial / pool modules -> | |
| its own `cwm3d_v2_nospatial_batonpre_pretrain_s123`; every other config shares the full trunk -> `cwm3d_v2_batonpre_pretrain_s123` | |
| (run with the surviving config's flags so the pretraining objective matches its fine-tune).""" | |
| return "cwm3d_v2_nospatial_batonpre_pretrain_s123" if cfg == "nospatial" else "cwm3d_v2_batonpre_pretrain_s123" | |
| def k_run(seed: int, pre: bool = False) -> str: | |
| return f"cwm3d_kp_batonpre_trunkp123_s{seed}_v2" if pre else f"cwm3d_kp_s{seed}_v2" | |
| def mb_changes_dl(cfg: str) -> bool: | |
| """True iff the MB grid point changes width / depth (-> a new pretraining `motionbert_kp_<tag>_batonpre_pretrain_s123` is needed).""" | |
| return MB_GRID[cfg][0] != "motionbert_kp" | |
| def mb_pretrain_run(cfg: str) -> str: | |
| return f"{MB_GRID[cfg][0]}_batonpre_pretrain_s123" | |
| def flags(family: str, cfg: str) -> list[str]: | |
| if family == "kv2": | |
| return ["--model", KV2_MODEL] + list(KV2_CONFIGS[cfg][0]) | |
| if family == "mb": | |
| model, extra = MB_GRID[cfg] | |
| return ["--model", model] + list(extra) | |
| raise ValueError(family) | |
| # ----------------------------------------------------------------------------- io | |
| def utc_now() -> str: | |
| return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) | |
| def run_json(name: str, runs: Path = RUNS): | |
| fs = sorted(glob.glob(str(Path(runs) / name / f"run_{name}_2*.json"))) | |
| return json.load(open(fs[-1])) if fs else None | |
| def finished(name: str, runs: Path = RUNS) -> bool: | |
| """A run counts as finished for the gate iff its run json (written after the single test evaluation) AND pred_val.npz exist.""" | |
| return (Path(runs) / name / "pred_val.npz").exists() and run_json(name, runs) is not None | |
| def val_data(): | |
| d = load_split("val") | |
| gt = d["S10"][:, P.T_OBS:] | |
| hm, thr, _ = M.hm_mask(gt) | |
| return d, gt, hm, thr | |
| def clip_errors(name: str, d: dict, gt: np.ndarray, runs: Path = RUNS, split: str = "val"): | |
| """[N,H,J] per-joint errors of a run's pred_<split>.npz aligned to the split's clip order; None if the file is missing.""" | |
| p = Path(runs) / name / f"pred_{split}.npz" | |
| if not p.exists(): | |
| return None | |
| z = np.load(p, allow_pickle=False) | |
| pred = np.asarray(z["pred"], dtype=np.float64)[align(z, d["clip_id"])] | |
| return M.joint_errors(pred, gt) | |
| def per_clip_all(err): | |
| return err.mean((1, 2)) | |
| def per_clip_hm(err, hm): | |
| return err[hm].mean((1, 2)) | |
| # ----------------------------------------------------------------------------- pure rules (unit-tested on synthetic arrays) | |
| def s1_row(cfg: str, run: str, cand_all, base_all, cand_hm, base_hm, n_components: int, n_boot: int = N_BOOT, seed: int = SEED) -> dict: | |
| """One S1 line from per-clip mean errors (All: n clips; HM: the val HM clips of both runs, same order).""" | |
| d_all = M.paired_diff(cand_all, base_all, n_boot, seed) | |
| d_hm = M.paired_diff(cand_hm, base_hm, n_boot, seed) | |
| survive_all = bool(d_all["mean"] <= S1_THRESH) | |
| guard_kill = bool(d_hm["mean"] > HM_GUARD and d_hm["ci_lo"] > 0) | |
| status = "SURVIVE" if survive_all and not guard_kill else ("GUARD_KILL" if survive_all else "KILL") | |
| return {"cfg": cfg, "run": run, "status": status, "n": int(len(cand_all)), "n_hm": int(len(cand_hm)), "n_components": int(n_components), | |
| "val_all": float(np.mean(cand_all)), "k_val_all": float(np.mean(base_all)), "val_hm": float(np.mean(cand_hm)), "k_val_hm": float(np.mean(base_hm)), | |
| "d_all": d_all, "d_hm": d_hm, "survive_all": survive_all, "guard_kill": guard_kill} | |
| def s1_select(rows: list[dict], max_survivors: int = S1_MAX_SURVIVORS, tie: float = S1_TIE) -> list[str]: | |
| """Ranked survivors (≤ max): best mean Δ first; a candidate within `tie` px of the current best with FEWER components wins the slot.""" | |
| remaining = sorted([r for r in rows if r.get("status") == "SURVIVE"], key=lambda r: (r["d_all"]["mean"], r["n_components"], r["cfg"])) | |
| out = [] | |
| while remaining and len(out) < max_survivors: | |
| best = remaining[0]["d_all"]["mean"] | |
| tied = [r for r in remaining if r["d_all"]["mean"] - best <= tie + 1e-12] | |
| pick = min(tied, key=lambda r: (r["n_components"], r["d_all"]["mean"], r["cfg"])) | |
| out.append(pick["cfg"]) | |
| remaining.remove(pick) | |
| return out | |
| def s1b_decide(d42_mean: float, d123_mean: float) -> tuple[bool, str]: | |
| two = 0.5 * (d42_mean + d123_mean) | |
| ok = bool(d123_mean < 0 and two <= S1B_THRESH) | |
| return ok, ("CONTINUE" if ok else "val-killed at seed 123") + f" (Δ val 42 {d42_mean:+.3f}, Δ val 123 {d123_mean:+.3f}, 2-seed mean {two:+.3f})" | |
| def s3_select(rows: list[dict], tie: float = S3_TIE) -> str | None: | |
| """rows: {cfg, val_all, n_params}; argmin val All, ties within `tie` -> fewer parameters, then lower val, then name.""" | |
| cands = [r for r in rows if r.get("val_all") is not None and r.get("n_params") is not None] | |
| if not cands: | |
| return None | |
| best = min(r["val_all"] for r in cands) | |
| tied = [r for r in cands if r["val_all"] - best <= tie + 1e-12] | |
| return min(tied, key=lambda r: (r["n_params"], r["val_all"], r["cfg"]))["cfg"] | |
| def dry_trunk_load(model, ckpt_path) -> dict: | |
| """Dry trunk-only load (CPU) + the reuse decision; `train_kp.load_init_ckpt` is the production loader (imported lazily).""" | |
| from aide.train.train_kp import load_init_ckpt | |
| prefixes = tuple(getattr(model, "TRUNK_PREFIXES", ("tok.", "blocks.", "ln_f."))) | |
| rep = load_init_ckpt(model, str(ckpt_path), "trunk", "cpu") | |
| trunk_keys = [k for k in model.state_dict() if k.startswith(prefixes)] | |
| n_missing_trunk = len(trunk_keys) - int(rep["n_loaded"]) | |
| reuse = bool(rep["n_skipped_shape"] == 0 and n_missing_trunk == 0) | |
| return {"ckpt": str(ckpt_path), "trunk_prefixes": list(prefixes), "n_trunk_tensors": len(trunk_keys), "n_loaded": int(rep["n_loaded"]), | |
| "n_skipped_shape": int(rep["n_skipped_shape"]), "n_missing_trunk": int(n_missing_trunk), "skipped_shape": rep.get("skipped_shape", []), | |
| "ckpt_epoch": rep.get("ckpt_epoch"), "ckpt_sha256": rep.get("ckpt_sha256"), "trunk_complete": reuse, | |
| "decision": "REUSE_K_CKPT" if reuse else "NEED_V2_PRETRAIN"} | |
| # ----------------------------------------------------------------------------- stages | |
| def _write(out: Path, stem: str, dec: dict) -> Path: | |
| out = Path(out) | |
| out.mkdir(parents=True, exist_ok=True) | |
| stamp = dec["utc"].replace("-", "").replace(":", "") | |
| (out / f"{stem}_{stamp}.json").write_text(json.dumps(dec, indent=1, default=str)) | |
| p = out / f"{stem}.json" | |
| p.write_text(json.dumps(dec, indent=1, default=str)) | |
| return p | |
| def stage_s1(runs: Path = RUNS, out: Path = KV2, n_boot: int = N_BOOT, seed: int = SEED, write: bool = True) -> dict: | |
| d, gt, hm, thr = val_data() | |
| base = k_run(42) | |
| base_err = clip_errors(base, d, gt, runs) | |
| if base_err is None: | |
| raise FileNotFoundError(f"baseline {base} has no pred_val.npz under {runs}") | |
| b_all, b_hm = per_clip_all(base_err), per_clip_hm(base_err, hm) | |
| rows = [] | |
| for cfg, (fl, ncomp) in KV2_CONFIGS.items(): | |
| name = kv2_run(cfg, 42) | |
| if not finished(name, runs): | |
| rows.append({"cfg": cfg, "run": name, "status": "PENDING", "n_components": ncomp, "flags": fl}) | |
| continue | |
| err = clip_errors(name, d, gt, runs) | |
| r = s1_row(cfg, name, per_clip_all(err), b_all, per_clip_hm(err, hm), b_hm, ncomp, n_boot, seed) | |
| r["flags"] = fl | |
| rj = run_json(name, runs) or {} | |
| r["n_params"], r["t_start_utc"] = rj.get("n_params"), rj.get("t_start_utc") | |
| rows.append(r) | |
| complete = all(r["status"] != "PENDING" for r in rows) | |
| survivors = s1_select(rows) | |
| dec = {"utc": utc_now(), "stage": "s1", "rule": RULE_S1, "baseline": base, "split": "val", "n": int(len(d["clip_id"])), "n_hm": int(hm.sum()), | |
| "hm_threshold_val_px": float(thr), "n_boot": n_boot, "seed": seed, "complete": complete, "rows": rows, "survivors": survivors, | |
| "note": "survivors are provisional while complete == false; downstream lanes require complete == true"} | |
| if write: | |
| _write(out, "s1_decision", dec) | |
| return dec | |
| def stage_s1b(cfg: str, runs: Path = RUNS, out: Path = KV2, n_boot: int = N_BOOT, seed: int = SEED, write: bool = True) -> dict: | |
| if cfg not in KV2_CONFIGS: | |
| raise ValueError(cfg) | |
| d, gt, hm, _ = val_data() | |
| rows, means = [], {} | |
| for s in (42, 123): | |
| name, base = kv2_run(cfg, s), k_run(s) | |
| if not (finished(name, runs) and (Path(runs) / base / "pred_val.npz").exists()): | |
| rows.append({"seed": s, "run": name, "baseline": base, "status": "PENDING"}) | |
| continue | |
| e_c, e_b = clip_errors(name, d, gt, runs), clip_errors(base, d, gt, runs) | |
| r = s1_row(cfg, name, per_clip_all(e_c), per_clip_all(e_b), per_clip_hm(e_c, hm), per_clip_hm(e_b, hm), KV2_CONFIGS[cfg][1], n_boot, seed) | |
| r.update(seed=s, baseline=base, status="DONE") | |
| rows.append(r) | |
| means[s] = r["d_all"]["mean"] | |
| complete = len(means) == 2 | |
| if complete: | |
| ok, reason = s1b_decide(means[42], means[123]) | |
| else: | |
| ok, reason = False, "PENDING (seed 42 and 123 runs needed)" | |
| dec = {"utc": utc_now(), "stage": "s1b", "cfg": cfg, "rule": RULE_S1B, "split": "val", "n_boot": n_boot, "seed": seed, "complete": complete, | |
| "rows": rows, "decision": "CONTINUE" if ok else ("KILLED" if complete else "PENDING"), "reason": reason, "survivors": [cfg] if ok else []} | |
| if write: | |
| _write(out, f"s1b_{cfg}_decision", dec) | |
| return dec | |
| def stage_s3(runs: Path = RUNS, out: Path = KV2, n_boot: int = N_BOOT, seed: int = SEED, write: bool = True) -> dict: | |
| d, gt, hm, _ = val_data() | |
| base = mb_run("default", 42) | |
| base_err = clip_errors(base, d, gt, runs) | |
| b_all = per_clip_all(base_err) if base_err is not None else None | |
| rows = [] | |
| for cfg, (model, fl) in MB_GRID.items(): | |
| name = mb_run(cfg, 42) | |
| if not finished(name, runs): | |
| rows.append({"cfg": cfg, "run": name, "model": model, "flags": fl, "status": "PENDING"}) | |
| continue | |
| err = clip_errors(name, d, gt, runs) | |
| rj = run_json(name, runs) or {} | |
| r = {"cfg": cfg, "run": name, "model": model, "flags": fl, "status": "DONE", "val_all": float(per_clip_all(err).mean()), | |
| "val_hm": float(per_clip_hm(err, hm).mean()), "n_params": rj.get("n_params"), "t_start_utc": rj.get("t_start_utc"), | |
| "d_all_vs_default": M.paired_diff(per_clip_all(err), b_all, n_boot, seed) if b_all is not None else None} | |
| rows.append(r) | |
| complete = all(r["status"] != "PENDING" for r in rows) | |
| best = s3_select([r for r in rows if r["status"] == "DONE"]) | |
| dec = {"utc": utc_now(), "stage": "s3", "rule": RULE_S3, "baseline": base, "split": "val", "n_boot": n_boot, "seed": seed, "complete": complete, "rows": rows, | |
| "mb_best": best, "mb_best_model": MB_GRID[best][0] if best else None, "mb_best_flags": MB_GRID[best][1] if best else None, | |
| "mb_best_is_default": (best == "default") if best else None, "mb_best_changes_dL": mb_changes_dl(best) if best else None, | |
| "mb_best_pretrain_run": mb_pretrain_run(best) if best else None, "survivors": [best] if best else [], | |
| "note": "mb_best is provisional while complete == false"} | |
| if write: | |
| _write(out, "s3_decision", dec) | |
| return dec | |
| def stage_pretrain_check(cfg: str = "full", ckpt: Path = K_CKPT, out: Path = KV2, write: bool = True, stem: str | None = None) -> dict: | |
| """Build the K-v2 model of `cfg` on CPU exactly as train_kp does (seed 123) and dry-load a pretraining checkpoint trunk-only (default: K's; | |
| `--ckpt <own v2 checkpoint> --stem pretrain_check_<cfg>_own` verifies the architecture-matched pretraining before the +pre fine-tunes: | |
| trunk_complete must be true = 0 shape-skipped and 0 missing trunk tensors).""" | |
| stem = stem or f"pretrain_check_{cfg}" | |
| import torch | |
| torch.set_num_threads(4) | |
| from aide.train import train_kp | |
| argv = ["--cpu", "--seed", "123"] + flags("kv2", cfg) | |
| try: | |
| a = train_kp.get_args(argv) | |
| train_kp.seed_all(123) | |
| model, _ = train_kp.build_model(a.model, a) | |
| except (SystemExit, ValueError, ImportError, ModuleNotFoundError, AttributeError) as e: # cwm3d_v2 not yet wired into train_kp | |
| dec = {"utc": utc_now(), "stage": "pretrain-check", "cfg": cfg, "rule": RULE_PRE, "status": "NOT_READY", "error": repr(e), "decision": None} | |
| if write: | |
| _write(out, stem, dec) | |
| return dec | |
| if not Path(ckpt).exists(): | |
| dec = {"utc": utc_now(), "stage": "pretrain-check", "cfg": cfg, "rule": RULE_PRE, "status": "NO_CKPT", "ckpt": str(ckpt), "decision": "NEED_V2_PRETRAIN"} | |
| else: | |
| rep = dry_trunk_load(model, ckpt) | |
| dec = {"utc": utc_now(), "stage": "pretrain-check", "cfg": cfg, "model": a.model, "trainer_argv": argv, "rule": RULE_PRE, "status": "DONE", | |
| "n_params": int(sum(p.numel() for p in model.parameters())), **rep} | |
| if write: | |
| _write(out, stem, dec) | |
| return dec | |
| # ----------------------------------------------------------------------------- cli | |
| def _print_summary(dec: dict): | |
| for r in dec.get("rows", []): | |
| if r.get("status") == "PENDING": | |
| print(f" {r.get('cfg', r.get('seed')):>14} PENDING ({r['run']})") | |
| elif "d_all" in r: | |
| da, dh = r["d_all"], r.get("d_hm") or {} | |
| print(f" {r['cfg']:>14} {r['status']:<10} val All {r['val_all']:.3f} Δ All {da['mean']:+.3f} [{da['ci_lo']:+.3f},{da['ci_hi']:+.3f}]" | |
| + (f" Δ HM {dh['mean']:+.2f} [{dh['ci_lo']:+.2f},{dh['ci_hi']:+.2f}]" if dh else "") + f" comps {r.get('n_components')}") | |
| elif "val_all" in r: | |
| dv = r.get("d_all_vs_default") or {} | |
| print(f" {r['cfg']:>14} {r['status']:<10} val All {r['val_all']:.3f} params {r.get('n_params')}" | |
| + (f" Δ vs default {dv['mean']:+.3f} [{dv['ci_lo']:+.3f},{dv['ci_hi']:+.3f}]" if dv else "")) | |
| def main(argv=None): | |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| ap.add_argument("stage", choices=("s1", "s1b", "s3", "pretrain-check", "flags", "name", "read")) | |
| ap.add_argument("args", nargs="*") | |
| ap.add_argument("--runs", default=str(RUNS)) | |
| ap.add_argument("--out", default=str(KV2)) | |
| ap.add_argument("--n-boot", type=int, default=N_BOOT) | |
| ap.add_argument("--ckpt", default=str(K_CKPT)) | |
| ap.add_argument("--stem", default=None, help="pretrain-check: json stem (default pretrain_check_<cfg>)") | |
| a = ap.parse_args(argv) | |
| runs, out = Path(a.runs), Path(a.out) | |
| if a.stage == "flags": | |
| print(" ".join(flags(a.args[0], a.args[1]))) | |
| return 0 | |
| if a.stage == "name" and a.args[0] == "kv2pre": | |
| print(kv2_pretrain_run(a.args[1])) | |
| return 0 | |
| if a.stage == "name": | |
| fam, cfg, seed = a.args[0], a.args[1], int(a.args[2]) | |
| pre = len(a.args) > 3 and a.args[3] == "pre" | |
| print(kv2_run(cfg, seed, pre) if fam == "kv2" else mb_run(cfg, seed, pre)) | |
| return 0 | |
| if a.stage == "read": | |
| p = out / f"{a.args[0]}.json" | |
| if not p.exists(): | |
| print("") | |
| return 3 | |
| v = json.load(open(p)).get(a.args[1]) | |
| print(" ".join(map(str, v)) if isinstance(v, list) else ("1" if v is True else "0" if v is False else "" if v is None else str(v))) | |
| return 0 | |
| if a.stage == "s1": | |
| dec = stage_s1(runs, out, a.n_boot) | |
| print(f"[kv2_gate s1 {dec['utc']}] complete={dec['complete']} survivors={dec['survivors']}") | |
| _print_summary(dec) | |
| return 0 if dec["complete"] else 3 | |
| if a.stage == "s1b": | |
| dec = stage_s1b(a.args[0], runs, out, a.n_boot) | |
| print(f"[kv2_gate s1b {dec['utc']}] cfg={dec['cfg']} complete={dec['complete']} decision={dec['decision']} — {dec['reason']}") | |
| return 0 if dec["complete"] else 3 | |
| if a.stage == "s3": | |
| dec = stage_s3(runs, out, a.n_boot) | |
| print(f"[kv2_gate s3 {dec['utc']}] complete={dec['complete']} mb_best={dec['mb_best']} (model {dec['mb_best_model']}, flags {dec['mb_best_flags']}, changes d/L: {dec['mb_best_changes_dL']})") | |
| _print_summary(dec) | |
| return 0 if dec["complete"] else 3 | |
| if a.stage == "pretrain-check": | |
| cfg = a.args[0] if a.args else "full" | |
| dec = stage_pretrain_check(cfg, Path(a.ckpt), out, stem=a.stem) | |
| print(f"[kv2_gate pretrain-check {dec['utc']}] cfg={cfg} ckpt={a.ckpt} status={dec['status']} decision={dec.get('decision')} trunk_complete={dec.get('trunk_complete')} " | |
| f"loaded={dec.get('n_loaded')} skipped_shape={dec.get('n_skipped_shape')} missing_trunk={dec.get('n_missing_trunk')} {dec.get('error', '')}") | |
| return 4 if dec["status"] == "NOT_READY" else 0 | |
| return 2 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |
Xet Storage Details
- Size:
- 23 kB
- Xet hash:
- 2ee869598da78d20af915a8f6f74a02791908335983eefb877c83490d1a903d4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.