"""Matched GPT-2 small eval for Stasis step 36,420 endpoint. Reuses the frozen package code
(stasis_gpt2_eval.py, METRIC_DEFINITIONS 579ab8c1...) unchanged; only the checkpoint loader differs
(accepts step-36420 endpoint by sha). Items, tokenization, tok_budget 8192, BOS, bootstrap identical.
Internal row key 'STASIS_12269' is the frozen Stasis-tokenizer key; outputs relabelled STASIS_36420."""
import json, os, sys, time, platform, argparse
import torch
PKG = os.environ.get("GPT2EVAL_PKG", "package")
sys.path.insert(0, os.path.join(PKG, "code")); sys.path.insert(0, os.path.join(PKG, "code"))
import stasis_gpt2_eval as E
SHA = "5b2dc755ed6e38817be84aa982e4591f64c8cc94ce1fa0316c2ac0336e0bcf85"
def load(ckpt, device):
    from frozen_stasis.config import ModelConfig
    from frozen_stasis.model import StasisDecoder
    E.need(E.file_sha(ckpt) == SHA, "step-36420 sha")
    p = torch.load(ckpt, map_location="cpu", weights_only=True, mmap=True)
    E.need(p["cursor"]["optimizer_step"] == 36420, "step 36420")
    sd = p["model"]; m = StasisDecoder(ModelConfig())
    E.need(list(sd) == list(m.state_dict()), "keys"); m.load_state_dict(sd, strict=True)
    E.need(sum(x.numel() for x in m.parameters()) == 109529856, "params")
    return E.StasisWrap(m).to(device).eval()
ap = argparse.ArgumentParser(); ap.add_argument("--ckpt"); ap.add_argument("--out"); ap.add_argument("--limit", type=int)
ap.add_argument("--gpt2", action="store_true", help="also run GPT-2 on this device"); ap.add_argument("--threads", type=int, default=8)
a = ap.parse_args(); torch.set_num_threads(a.threads); os.makedirs(a.out, exist_ok=True)
t0 = time.time(); lf = open(os.path.join(a.out, "run_log.txt"), "a")
def log(m):
    s = f"[{time.time()-t0:7.1f}s] {m}"; print(s, flush=True); lf.write(s + "\n"); lf.flush()
torch.manual_seed(20261016); sums = E.verify_package(PKG); log(f"package verified {len(sums)} files metric {E.METRIC_SHA[:16]}")
if a.gpt2:
    g = E.load_gpt2(PKG, "cpu"); rg = E.run_model("GPT2_SMALL", g, PKG, "cpu", 8192, a.limit, a.out, log); del g
s = load(a.ckpt, "cpu"); log("Stasis step-36420 loaded strict weights_only")
rs = E.run_model("STASIS_12269", s, PKG, "cpu", 8192, a.limit, a.out, log)
for b in E.BENCHES:
    os.replace(os.path.join(a.out, f"ROWS_{b}_STASIS_12269.jsonl"), os.path.join(a.out, f"ROWS_{b}_STASIS_36420.jsonl"))
log("done")