fix codex BLOCKERs: global fixed layout invariance, min_steps off-by-one, integers EOS alias, rerun guard, eval prereg-literal codes + dual-checkpoint honesty; +7 regression tests

1 files changed, 107 insertions, 64 deletions

diff --git a/src/eval.py b/src/eval.py
index 4485523..fdc2dec 100644
--- a/src/eval.py
+++ b/src/eval.py
@@ -1,7 +1,9 @@
 """Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature.
 
-Interpretation codes are locked in design/preregistration.md — this module only MEASURES
-and classifies against those definitions (O1-O4, H1-H4, P1-P4).
+Interpretation codes are locked in design/preregistration.md (with operationalization
+thresholds in Addendum 2). This module MEASURES and classifies against those definitions.
+Where the prereg prose supplies no testable boundary, the code emits the measurements
+plus "unclassified" rather than inventing a label.
 """
 import argparse
 import csv
@@ -23,38 +25,44 @@ def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
     primes = sieve_primes(hi + 200)
     correct = 0
     errors = []
-    easy_misses = 0
+    easy_total = 0
+    easy_wrong = 0
     for n in range(lo, hi + 1):
         x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
         gen = greedy_decode(model, x, cfg)[0].tolist()
         pred = decode_tokens(gen, cfg)
         target = next_prime(n, primes)
+        is_easy = (n % 2 == 0) or (n % 5 == 0)   # trivial composites (skip-evens / skip-5s)
+        if is_easy:
+            easy_total += 1
         if pred == target:
             correct += 1
         else:
             errors.append({"n": n, "target": target, "pred": pred})
-            if n % 2 == 0 or n % 5 == 0:
-                easy_misses += 1
+            if is_easy:
+                easy_wrong += 1
     total = hi - lo + 1
     acc = correct / total
     flagged = [e for e in errors if e["n"] in FLAGGED_COMPOSITES]
-    # P-code classification (see preregistration.md)
-    if acc >= 0.85:
-        code = "P3"
-    elif easy_misses >= 5:
-        code = "P4"
-    elif errors and all(e["n"] in FLAGGED_COMPOSITES for e in errors) and len(errors) <= 6:
-        code = "P1"
+    # classification per prereg + Addendum 2 operationalization; P4 checked first
+    if easy_total and easy_wrong / easy_total > 0.5:
+        code = "P4"      # fails trivial evens/5-multiples -> pure memorization
+    elif acc >= 0.85:
+        code = "P3"      # surprising success beyond expectation
+    elif len(flagged) >= 3 and len(errors) <= 6 and all(e["n"] in FLAGGED_COMPOSITES for e in errors):
+        code = "P1"      # errors concentrated on composites needing divisors 11,13 -> learned sieve
     else:
-        code = "P2"
+        code = "P2"      # scattered errors -> memorization / non-transferable heuristics
     return {
         "code": code, "acc": acc, "correct": correct, "total": total,
-        "errors": errors, "flagged_errors": flagged, "easy_misses": easy_misses,
+        "errors": errors, "flagged_errors": flagged,
+        "easy_total": easy_total, "easy_wrong": easy_wrong,
+        "easy_err_rate": (easy_wrong / easy_total) if easy_total else None,
     }
 
 
 def grokking_signature(metrics_path: str) -> dict:
-    """Classify the training curve against preregistered codes O1-O4."""
+    """Classify the training curve against preregistered codes O1-O4 (Addendum 2 operationalization)."""
     with open(metrics_path) as fh:
         rows = list(csv.DictReader(fh))
     if not rows:
@@ -62,31 +70,41 @@ def grokking_signature(metrics_path: str) -> dict:
     train = [float(r["train_em"]) for r in rows]
     val = [float(r["val_em"]) for r in rows]
     n = len(rows)
-    saturated = any(all(t >= 0.95 for t in train[i:i + 10]) for i in range(n - 9)) if n >= 10 else False
+    # first index where train EM >= 0.95 for 10 CONSECUTIVE evals
+    sat_start = next((i for i in range(n - 9) if all(t >= 0.95 for t in train[i:i + 10])), None)
     hi = next((i for i, v in enumerate(val) if v >= 0.9), None)
-    trans = None
-    if hi is not None:
-        lo_cands = [i for i in range(hi) if val[i] <= 0.2]
+    # transition: last eval with val <= 0.2 strictly before hi, and after saturation window
+    lo = None
+    if hi is not None and sat_start is not None:
+        lo_cands = [i for i in range(sat_start + 10, hi) if val[i] <= 0.2]
         if lo_cands:
-            trans = hi - max(lo_cands)
-    if saturated and hi is not None and trans is not None and trans <= 5:
-        code = "O1"
-    elif max(train) >= 0.95 and hi is not None and (trans is None or trans > 5):
-        code = "O3"
-    elif max(train) >= 0.95 and hi is None:
-        code = "O2"
+            lo = max(lo_cands)
+    width = (hi - lo) if (hi is not None and lo is not None) else None
+
+    max_train = max(train)
+    if max_train < 0.95:
+        code = "O4"          # train never reached 0.95 -> setup/optimization failure
+    elif sat_start is None:
+        code = "O4"          # reached 0.95 but never sustained 10 evals within budget
+    elif hi is not None and width is not None and width <= 5:
+        code = "O1"          # sharp transition AFTER sustained train saturation
+    elif hi is not None:
+        code = "O3"          # reached 0.9+ but not via the sharp O1 pattern (gradual)
+    elif val[-1] <= 0.3:
+        code = "O2"          # train memorized, val stayed low
     else:
-        code = "O4"
+        code = "O-PARTIAL"   # val ended in (0.3, 0.9) with no 0.9 reach — prereg has no boundary
     return {
-        "code": code, "train_saturated": saturated, "val_hi_eval_idx": hi,
-        "transition_width_evals": trans, "n_evals": n,
+        "code": code, "sat_start_eval": sat_start, "val_hi_eval_idx": hi,
+        "transition_width_evals": width, "n_evals": n,
         "final_train_em": train[-1], "final_val_em": val[-1],
+        "max_train_em": max_train,
     }
 
 
 @torch.no_grad()
 def halting_report(model, cfg: Config) -> dict:
-    """RNN halting structure: mean steps + correlation with gap-to-next-prime (H1-H4)."""
+    """RNN halting structure: mean steps at run end + correlation with gap-to-next-prime (H1-H4)."""
     primes = sieve_primes(300)
     gaps, steps = [], []
     for n in range(cfg.range_start, cfg.range_end + 1):
@@ -97,54 +115,79 @@ def halting_report(model, cfg: Config) -> dict:
         steps.append(float(s.mean()))
     mean = float(np.mean(steps))
     rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0
-    if mean <= cfg.min_steps + 0.5:
-        code = "H1"
-    elif mean >= cfg.max_steps - 0.5:
-        code = "H2"
-    elif abs(rho) >= 0.3:
-        code = "H3"
+    lo, hi = cfg.min_steps + 0.5, cfg.max_steps - 0.5
+    if mean <= lo:
+        code = "H1"          # collapse to the min_steps floor
+    elif mean >= hi:
+        code = "H2"          # pinned at K — never learned to halt
+    elif rho >= 0.3:
+        code = "H3"          # intermediate + positive gap correlation -> computation budget
     else:
-        code = "H4"
-    return {"code": code, "mean_steps": mean, "corr_gap": rho, "min_steps": cfg.min_steps, "max_steps": cfg.max_steps}
+        code = "H4"          # intermediate but no positive gap correlation — noisy/unused
+    return {"code": code, "mean_steps": mean, "corr_gap": rho,
+            "min_steps": cfg.min_steps, "max_steps": cfg.max_steps}
+
+
+def _load(out_dir: str, ckpt: str, cfg: Config):
+    model = build_model(cfg)
+    model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location="cpu"))
+    model.eval()
+    return model
 
 
 def main() -> None:
     ap = argparse.ArgumentParser(description="prime-grokking eval")
     ap.add_argument("model")
     ap.add_argument("seed")
-    ap.add_argument("--ckpt", default="best.pt")
     a = ap.parse_args()
     out_dir = os.path.join("runs", a.model, f"seed{a.seed}")
     cfg = Config.load(os.path.join(out_dir, "config.json"))
-    model = build_model(cfg)
-    model.load_state_dict(torch.load(os.path.join(out_dir, a.ckpt), map_location="cpu"))
-    model.eval()
 
     _, val_in = get_splits(cfg)
     val_ex = build_examples(val_in, cfg)
-    tok, em, per = evaluate(model, val_ex, cfg)
-
-    probe = probe_report(model, cfg)
-    sig = grokking_signature(os.path.join(out_dir, "metrics.csv"))
-    hlt = halting_report(model, cfg) if cfg.model == "rnn" else None
-
-    results = {
-        "model": cfg.model, "seed": cfg.seed, "ckpt": a.ckpt,
-        "params": model.param_count(),
-        "val_token_acc": tok, "val_exact_match": em,
-        "per_example": [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
-                        for x, t, p, ok in per],
-        "probe": probe,
-        "signature": sig,
-        "halting": hlt,
-    }
+
+    report = {"model": cfg.model, "seed": cfg.seed, "params": None, "val_selected": {}, "final": {}}
+    for ckpt in ("best.pt", "last.pt"):
+        path = os.path.join(out_dir, ckpt)
+        if not os.path.exists(path):
+            continue
+        model = _load(out_dir, ckpt, cfg)
+        tok, em, per = evaluate(model, val_ex, cfg)
+        entry = {
+            "ckpt": ckpt,
+            "params": model.param_count(),
+            "val_token_acc": tok,
+            "val_exact_match": em,
+            "note": ("val-selected checkpoint (early stop / best-on-val per prereg — "
+                     "val EM here is selection-holed, not an untouched test estimate" if ckpt == "best.pt"
+                     else "final checkpoint, no val selection"),
+        }
+        if ckpt == "last.pt":
+            entry["probe"] = probe_report(model, cfg)
+            if cfg.model == "rnn":
+                entry["halting"] = halting_report(model, cfg)
+            entry["per_example"] = [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
+                                    for x, t, p, ok in per]
+        if ckpt == "best.pt":
+            report["val_selected"] = entry
+        else:
+            report["final"] = entry
+
+    report["signature"] = grokking_signature(os.path.join(out_dir, "metrics.csv"))
     with open(os.path.join(out_dir, "results.json"), "w") as fh:
-        json.dump(results, fh, indent=2)
-    print(json.dumps({"val_token_acc": round(tok, 4), "val_em": round(em, 4),
-                      "probe": probe["code"], "probe_acc": round(probe["acc"], 3),
-                      "flagged_errors": [e["n"] for e in probe["flagged_errors"]],
-                      "signature": sig["code"], "halting": hlt["code"] if hlt else None,
-                      "results": os.path.join(out_dir, "results.json")}, indent=2))
+        json.dump(report, fh, indent=2)
+
+    f = report["final"]
+    p = f.get("probe", {})
+    print(json.dumps({
+        "val_em_best": round(report["val_selected"].get("val_exact_match", float("nan")), 4),
+        "val_em_last": round(f.get("val_exact_match", float("nan")), 4),
+        "probe_code": p.get("code"), "probe_acc": round(p.get("acc", float("nan")), 3),
+        "flagged_errors": [e["n"] for e in p.get("flagged_errors", [])],
+        "signature": report["signature"]["code"],
+        "halting": (f.get("halting") or {}).get("code"),
+        "results": os.path.join(out_dir, "results.json"),
+    }, indent=2))
 
 
 if __name__ == "__main__":