File size: 2,180 Bytes
093306a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """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
|