| """Entry point shared by the Kaggle notebook and the box CPU smoke test.""" |
| import json, os, platform, sys, time |
| import torch |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import stasis_gpt2_eval as E |
|
|
| def main(pkg, ckpt, outdir, device, limit=None, tok_budget=8192, ckpt_sha=None, label="REAL"): |
| os.makedirs(outdir, exist_ok=True); t0 = time.time(); logf = open(os.path.join(outdir, "run_log.txt"), "a") |
| def log(m): |
| s = f"[{time.time()-t0:7.1f}s] {m}"; print(s, flush=True); logf.write(s + "\n"); logf.flush() |
| torch.backends.cuda.matmul.allow_tf32 = False; torch.backends.cudnn.allow_tf32 = False; torch.manual_seed(20261016) |
| sums = E.verify_package(pkg); log(f"package verified: {len(sums)} files, metric sha {E.METRIC_SHA[:16]}") |
| env = {"torch": torch.__version__, "python": platform.python_version(), "device": device, |
| "gpu": torch.cuda.get_device_name(0) if device.startswith("cuda") else None, "limit": limit, "tok_budget": tok_budget} |
| g = E.load_gpt2(pkg, device); log("GPT-2 loaded (strict)") |
| rg = E.run_model("GPT2_SMALL", g, pkg, device, tok_budget, limit, outdir, log); del g |
| if device.startswith("cuda"): torch.cuda.empty_cache() |
| s, sha = E.load_stasis(pkg, ckpt, device, ckpt_sha); log(f"Stasis loaded (strict, weights_only) sha {sha[:16]}") |
| rs = E.run_model("STASIS_12269", s, pkg, device, tok_budget, limit, outdir, log); del s |
| summary = E.summarize(rs, rg); log("summary computed") |
| cont = json.load(open(os.path.join(pkg, "CONTAMINATION_13GRAM.json"))) |
| result = {"label": label, "metric_definitions_sha256": E.METRIC_SHA, "checkpoint": {"path": ckpt, "sha256": sha, "file": E.CKPT_SHAS[sha]}, |
| "gpt2_safetensors_sha256": sums["gpt2/model.safetensors"], "env": env, "elapsed_s": time.time() - t0, |
| "results": summary, "contamination_13gram_sixth_tranche": cont, |
| "reference_step3403": {"panel": "WT2_PANEL256_CONTINUITY", "stasis_bpb": 1.6061006758479426, "gpt2_bpb": 1.31791335174742, "deficit_pct": 21.86694}} |
| json.dump(result, open(os.path.join(outdir, "RESULTS.json"), "w"), indent=1); log("wrote RESULTS.json") |
| return result |
|
|
