1#!/usr/bin/env python2"""Locked evaluation driver. Reads an evaluation YAML listing systems (checkpoints + policies) and splits, refuses3to read any split whose name starts with `test` unless protocol/locked_test_manifest.json has "unsealed": true,4and writes one JSON per (system, split) under results/raw/ with full provenance.56Usage: python -m evaluation.run --config configs/experiments/locked_evaluation.yaml [--device mps] [--only NAME]7"""8from __future__ import annotations910import argparse, json, sys, time11from pathlib import Path1213import numpy as np14import torch1516sys.path.insert(0, str(Path(__file__).resolve().parents[1]))17from prisme.flops import action_cost_proxy # noqa: E40218from training.common import TokenBin, load_checkpoint, load_yaml, pick_device, seed_all, git_sha, sha256_file # noqa: E40219from training.train_editor import build_model, sample_forced_actions # noqa: E40220from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, exact_match, load_items # noqa: E40221from evaluation.perplexity import perplexity_by_policy # noqa: E402222324def check_lock(split: str):25 if split.startswith("test"):26 man = Path("protocol/locked_test_manifest.json")27 ok = man.exists() and json.loads(man.read_text()).get("unsealed", False)28 if not ok:29 raise SystemExit(f"REFUSED: split '{split}' is locked (protocol/locked_test_manifest.json not unsealed)")303132def main(argv=None) -> int:33 ap = argparse.ArgumentParser(description=__doc__)34 ap.add_argument("--config", required=True)35 ap.add_argument("--device", default="auto")36 ap.add_argument("--only", default=None)37 args = ap.parse_args(argv)38 exp = load_yaml(args.config); device = pick_device(args.device)39 out_dir = Path(exp.get("out_dir", "results/raw")); out_dir.mkdir(parents=True, exist_ok=True)40 for sysname, sysspec in exp["systems"].items():41 if args.only and sysname != args.only:42 continue43 seed_all(sysspec.get("seed", 1))44 cfg, model = build_model(sysspec); model = model.to(device).eval()45 meta = load_checkpoint(sysspec["checkpoint"], model, strict=False)46 for split in exp["splits"]:47 check_lock(split)48 t0 = time.time()49 data = exp["data"]50 rec = {"system": sysname, "split": split, "checkpoint": sysspec["checkpoint"], "checkpoint_sha256": meta.get("weights_sha256"),51 "config_sha256": cfg.sha256(), "code_commit": git_sha(), "device": str(device), "precision": "fp32", "evaluated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z")}52 if data["kind"] == "tokens":53 tb = TokenBin(Path(data["dir"]) / f"{split}.bin", exp["seq_len"])54 rec["data_sha256"] = sha256_file(Path(data["dir"]) / "manifest.json")55 cost = torch.tensor(action_cost_proxy(cfg, exp["seq_len"]), device=device)56 rec["perplexity"] = perplexity_by_policy(model, cfg, lambda: tb.sequential_batches(exp["batch_size"], device, exp.get("n_batches")),57 device, cost, budgets=exp.get("budgets", [0, 1, 2, 4, 8]), controller=sysspec.get("controller", True))58 else:59 tok = AnyTokenizer(data.get("tokenizer", "char"))60 items = load_items(Path(data["splits_dir"]) / f"{split}.jsonl")61 rec["data_sha256"] = sha256_file(Path(data["splits_dir"]) / "manifest.json")62 sb = SyntheticBatches(items, tok, exp["seq_len"])63 cost = torch.tensor(action_cost_proxy(cfg, exp["seq_len"]), device=device)64 rec["exact_match"] = {}65 policies = sysspec.get("policies", ["budget0"])66 for pol in policies:67 rng = np.random.default_rng(sysspec.get("seed", 1) + 7)68 if pol == "budget0":69 em = exact_match(model, sb, tok, device, 0, batch_size=exp["batch_size"], max_items=exp.get("max_items"))70 elif pol.startswith("forced"):71 b = int(pol.split("_b")[1])72 em = exact_match(model, sb, tok, device, b, batch_size=exp["batch_size"], max_items=exp.get("max_items"),73 forced_actions_fn=lambda B_, n, b=b: sample_forced_actions(cfg, B_, n, [b], [1.0], rng, device)[0])74 elif pol == "controller":75 em = exact_match(model, sb, tok, device, cfg.r_max, cost=cost, batch_size=exp["batch_size"], max_items=exp.get("max_items"))76 elif pol.startswith("controller_lam"):77 lam = float(pol.split("lam")[1])78 em = exact_match(model, sb, tok, device, cfg.r_max, cost=cost, batch_size=exp["batch_size"], max_items=exp.get("max_items"), lambda_cost=lam)79 else:80 raise ValueError(pol)81 recs = em.pop("records")82 rec["exact_match"][pol] = em83 (out_dir / f"items_{sysname}_{split}_{pol}.json").write_text(json.dumps(recs))84 print(f"{sysname} {split} {pol}: acc {em['accuracy']:.4f} n={em['n']} edits/tok {em['edits_per_token']:.2f}", flush=True)85 rec["seconds"] = round(time.time() - t0, 1)86 (out_dir / f"eval_{sysname}_{split}.json").write_text(json.dumps(rec, indent=1))87 print("wrote", out_dir / f"eval_{sysname}_{split}.json", flush=True)88 return 0899091if __name__ == "__main__":92 raise SystemExit(main())93