File size: 1,929 Bytes
f930dac | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | #!/usr/bin/env python3
"""Load every published checkpoint and check it against its recorded config.
python3 eval/verify_checkpoints.py # all runs
python3 eval/verify_checkpoints.py --sha # also verify SHA-256
Exits non-zero if any checkpoint fails to load or disagrees with its record.
"""
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))
from pns.checkpoint import verify # noqa: E402
from pns.common import ckpt_root, sha256_file # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--sha", action="store_true", help="also verify SHA-256")
ap.add_argument("--device", default="cpu")
args = ap.parse_args()
index = json.loads((ckpt_root() / "index.json").read_text())
sums = {}
p = ROOT / "results" / "artifact_hashes.csv"
if args.sha and p.exists():
import csv
for row in csv.DictReader(p.open()):
sums[row["path"]] = row["sha256"]
bad = []
for run in sorted(index):
r = verify(run, args.device)
ok = r["loaded"] and r["tensor_total_match"] and r["params_match"]
status = "ok" if ok else "MISMATCH"
line = f" {run:18s} {r['kind']:8s} {r['params']:>12,} params {status}"
if args.sha:
rel = f"checkpoints/{run}/model.safetensors"
want = sums.get(rel)
got = sha256_file(ckpt_root() / run / "model.safetensors")
shaok = want is None or want == got
ok = ok and shaok
line += f" sha {'ok' if shaok else 'MISMATCH'}"
print(line, flush=True)
if not ok:
bad.append(run)
print(f"\n{len(index) - len(bad)}/{len(index)} checkpoints verified")
if bad:
print("FAILED:", ", ".join(bad), file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()
|