1#!/usr/bin/env python2"""Routing comparison on a split (teacher-forced loss + extra ops) for: frozen backbone (budget 0), random /3uniform / fixed-pair schedules at budgets {1,2,4,8}, entropy and norm routers (threshold grid), the learned4controller (locked λ, κ; and a λ grid for the Pareto curve), the greedy one-step oracle and (on a small subset)5the two-step oracle. Oracles are DIAGNOSTIC (they read the target).67Usage: python -m evaluation.routing --config configs/experiments/routing_p4.yaml --split dev_iid [--device mps]8Writes results/raw/routing_<run>.json (one record per policy with loss, edits/token, extra ops ratio).9"""10from __future__ import annotations1112import argparse, json, sys, time13from pathlib import Path1415import numpy as np16import torch1718sys.path.insert(0, str(Path(__file__).resolve().parents[1]))19from prisme.flops import action_cost_proxy # noqa: E40220from prisme.losses import token_nll # noqa: E40221from prisme.oracle import greedy_oracle_rollout, two_step_oracle # noqa: E40222from baselines.routers import rollout as router_rollout # noqa: E40223from training.common import TokenBin, load_checkpoint, load_yaml, pick_device, seed_all, git_sha # noqa: E40224from training.train_editor import build_model, sample_forced_actions # noqa: E40225from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, load_items # noqa: E402262728def ops_of(actions, cost, cfg):29 a = actions30 c = torch.cat([cost, torch.zeros(1, device=cost.device)])31 return c[a.clamp(max=cfg.n_actions)].sum(-1)323334@torch.no_grad()35def evaluate_policy(model, cfg, batches, device, cost, kind, **kw):36 tot, n, ed, ops = 0.0, 0, 0.0, 0.037 extra = {}38 for x, y in batches:39 v = y != -10040 if kind == "budget0":41 logits = model(x, None)["logits"]; acts = None42 elif kind == "forced":43 rng = kw["rng"]; acts, _ = sample_forced_actions(cfg, x.shape[0], x.shape[1], [kw["budget"]], [1.0], rng, device)44 out = model(x, acts); logits = out["logits"]45 elif kind in ("uniform", "fixed_pair", "entropy", "norm", "random"):46 out = router_rollout(model, x, kind, kw["budget"], threshold=kw.get("threshold", 0.0), seed=kw.get("seed", 0))47 logits, acts = out["logits"], out["actions"]48 elif kind == "controller":49 out = model.forward_policy(x, kw["budget"], cost, lambda_cost=kw.get("lam"), kappa_risk=kw.get("kap"))50 logits, acts = out["logits"], out["actions"]51 elif kind == "oracle1":52 out = greedy_oracle_rollout(model, x, y, kw["budget"])53 nll = out["loss_per_step"][-1]; acts = out["actions"]; logits = None54 else:55 raise ValueError(kind)56 if logits is not None:57 nll = token_nll(logits, y)58 tot += float(nll[v].sum()); n += int(v.sum())59 if acts is not None:60 ed += float((acts != cfg.stop_action).sum(-1)[v].float().sum()); ops += float(ops_of(acts, cost, cfg)[v].sum())61 return {"loss": tot / max(n, 1), "edits_per_token": ed / max(n, 1), "extra_ops_ratio": ops / max(n, 1), "n_tokens": n}626364def main(argv=None) -> int:65 ap = argparse.ArgumentParser(description=__doc__)66 ap.add_argument("--config", required=True)67 ap.add_argument("--split", default=None)68 ap.add_argument("--device", default="auto")69 ap.add_argument("--out", default=None)70 args = ap.parse_args(argv)71 exp = load_yaml(args.config); seed_all(exp.get("seed", 1)); device = pick_device(args.device)72 cfg, model = build_model(exp); model = model.to(device).eval()73 meta = load_checkpoint(exp["checkpoint"], model, strict=False)74 data = exp["data"]; split = args.split or data.get("split", "dev_iid")75 bs, T = exp["batch_size"], exp["seq_len"]76 def batches(n=None):77 if data["kind"] == "tokens":78 return TokenBin(data["bin"], T).sequential_batches(bs, device, n or exp.get("n_batches"))79 tok = AnyTokenizer(data.get("tokenizer", "char"))80 sb = SyntheticBatches(load_items(Path(data["splits_dir"]) / f"{split}.jsonl"), tok, T)81 return sb.sequential_batches(bs, device, data.get("loss_span", "answer"), (n or exp.get("n_batches", 10**9)) * bs)82 cost = torch.tensor(action_cost_proxy(cfg, T), device=device)83 budgets = exp.get("budgets", [1, 2, 4, 8])84 recs = []85 def add(name, res, **tags):86 r = {"policy": name, **tags, **res}; recs.append(r); print(json.dumps(r), flush=True)87 t0 = time.time()88 add("budget0", evaluate_policy(model, cfg, batches(), device, cost, "budget0"))89 for b in budgets:90 add("random_forced", evaluate_policy(model, cfg, batches(), device, cost, "forced", budget=b, rng=np.random.default_rng(0)), budget=b)91 add("uniform", evaluate_policy(model, cfg, batches(), device, cost, "uniform", budget=b), budget=b)92 add("fixed_pair", evaluate_policy(model, cfg, batches(), device, cost, "fixed_pair", budget=b), budget=b)93 for thr in exp.get("entropy_thresholds", [0.1, 0.5, 1.0, 2.0]):94 add("entropy", evaluate_policy(model, cfg, batches(), device, cost, "entropy", budget=max(budgets), threshold=thr), threshold=thr, budget=max(budgets))95 for thr in exp.get("norm_thresholds", [0.02, 0.05, 0.1, 0.2]):96 add("norm", evaluate_policy(model, cfg, batches(), device, cost, "norm", budget=max(budgets), threshold=thr), threshold=thr, budget=max(budgets))97 if exp.get("controller", True):98 add("controller_locked", evaluate_policy(model, cfg, batches(), device, cost, "controller", budget=max(budgets), lam=cfg.lambda_cost, kap=cfg.kappa_risk), lam=cfg.lambda_cost, kap=cfg.kappa_risk, budget=max(budgets))99 for lam in exp.get("lambda_grid", [0.0, 0.01, 0.03, 0.1, 0.3]):100 add("controller", evaluate_policy(model, cfg, batches(), device, cost, "controller", budget=max(budgets), lam=lam, kap=cfg.kappa_risk), lam=lam, kap=cfg.kappa_risk, budget=max(budgets))101 nb_or = exp.get("oracle_batches", 4)102 for b in budgets:103 add("oracle1_greedy", evaluate_policy(model, cfg, batches(nb_or), device, cost, "oracle1", budget=b), budget=b, diagnostic=True)104 # two-step oracle on a small subset105 nb2 = exp.get("two_step_batches", 1)106 acc = {"one": 0.0, "two": 0.0, "greedy_then_best": 0.0, "n": 0}107 for i, (x, y) in enumerate(batches(nb2)):108 o = two_step_oracle(model, x, y)109 v = y != -100110 acc["one"] += float(o["one_step_best"][v].sum()); acc["two"] += float(o["two_step_best"][v].sum())111 acc["greedy_then_best"] += float(o["greedy_then_best"][v].sum()); acc["n"] += int(v.sum())112 add("oracle_two_step", {"one_step_gain": acc["one"] / max(acc["n"], 1), "two_step_gain": acc["two"] / max(acc["n"], 1),113 "greedy_then_best_gain": acc["greedy_then_best"] / max(acc["n"], 1), "n_tokens": acc["n"]}, diagnostic=True)114 out = {"config": args.config, "checkpoint": exp["checkpoint"], "checkpoint_sha256": meta.get("weights_sha256"), "split": split,115 "code_commit": git_sha(), "device": str(device), "seconds": round(time.time() - t0, 1), "records": recs}116 path = Path(args.out or f"results/raw/routing_{Path(exp['checkpoint']).parent.name}_{split}.json")117 path.parent.mkdir(parents=True, exist_ok=True); path.write_text(json.dumps(out, indent=1))118 print("wrote", path)119 return 0120121122if __name__ == "__main__":123 raise SystemExit(main())124