1#!/usr/bin/env python2"""Reproduce / update the analytic ops-per-token table of the WP31 brief (Sec. 3.9) from the implemented model.3Usage: python scripts/estimate_flops.py --config configs/prisme_459m.yaml --T 4096 [--json results/tables/ops_459m.json]4"""5from __future__ import annotations67import argparse, json, sys8from pathlib import Path910sys.path.insert(0, str(Path(__file__).resolve().parents[1]))11from prisme.config import load_config # noqa: E40212from prisme.flops import ops_per_token, action_cost_proxy # noqa: E4021314BRIEF_459M_T4096 = {"backbone_only": 1.296e9, "4_actions_k1": 1.645e9, "4_actions_k2": 1.767e9,15 "8_actions_k2": 2.195e9, "8_actions_dense16": 5.608e9}161718def main() -> int:19 ap = argparse.ArgumentParser(description=__doc__)20 ap.add_argument("--config", default="configs/prisme_459m.yaml")21 ap.add_argument("--T", type=int, default=4096)22 ap.add_argument("--json", default=None)23 args = ap.parse_args()24 c = load_config(args.config)25 rows = {26 "backbone_only": ops_per_token(c, args.T, 0, 1),27 "4_actions_k1": ops_per_token(c, args.T, 4, 1),28 "4_actions_k2": ops_per_token(c, args.T, 4, 2),29 "8_actions_k2": ops_per_token(c, args.T, 8, 2),30 "8_actions_dense16": ops_per_token(c, args.T, 8, c.n_slots),31 }32 out = {"config": c.name, "T": args.T, "assumptions": __import__("prisme.flops", fromlist=["x"]).__doc__,33 "rows": {k: v.as_dict() for k, v in rows.items()}, "brief_reference_459m_T4096": BRIEF_459M_T4096,34 "action_cost_proxy": action_cost_proxy(c, args.T)}35 print(f"{c.name} @ T={args.T}: ops/token (x1e9); controller evaluated n_actions+1 times (capped at r_max)")36 print(f"{'policy':20s} {'total':>8s} {'backbone':>9s} {'workspace':>10s} {'editor':>8s} {'ctrl':>7s} {'readout':>8s} {'brief':>8s}")37 for k, v in rows.items():38 ed = (v.editor_source_attn + v.editor_slot_attn + v.editor_ffn + v.editor_gate) / 1e939 ref = BRIEF_459M_T4096[k] / 1e9 if (c.name == "prisme_459m" and args.T == 4096) else float("nan")40 print(f"{k:20s} {v.total/1e9:8.3f} {v.backbone/1e9:9.3f} {v.workspace/1e9:10.3f} {ed:8.3f} {v.controller/1e9:7.3f} {v.readout/1e9:8.3f} {ref:8.3f}")41 print("cost proxy c(a) (single, pair):", round(out["action_cost_proxy"][0], 4), round(out["action_cost_proxy"][-1], 4))42 if args.json:43 Path(args.json).parent.mkdir(parents=True, exist_ok=True)44 Path(args.json).write_text(json.dumps(out, indent=2))45 return 0464748if __name__ == "__main__":49 raise SystemExit(main())50