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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
|
"""Classification correctness for src/eval.py — the preregistered O/H/P codes.
Ground truth is synthetic: stub models with KNOWN behavior feed the classifier,
and synthetic metrics CSVs feed grokking_signature. (Original: ad-hoc verifier
/tmp/hermes-verify-eval.py, promoted to a permanent regression suite.)
"""
import csv
import os
import tempfile
import torch
from src.config import Config
from src.data import sieve_primes
from src.eval import FLAGGED_SIEVE_PREDS, grokking_signature, halting_report, probe_report
DIGITS_CFG = Config(vocab_mode="digits")
EOS = DIGITS_CFG.eos_id
PAD = DIGITS_CFG.pad_id
def _decode_row(row) -> int:
toks = [int(t) for t in row.tolist() if int(t) != PAD]
return int("".join(map(str, toks))) if toks else -1
class StubModel(torch.nn.Module):
"""One-hot logits forcing greedy_decode to emit exactly pred_fn(n)."""
def __init__(self, pred_fn):
super().__init__()
self.pred_fn = pred_fn
def forward(self, x, y_in):
B, T_out = x.shape[0], y_in.shape[1]
logits = torch.full((B, T_out, 11), -100.0)
for i in range(B):
n = _decode_row(x[i])
toks = [int(d) for d in str(self.pred_fn(n))] + [EOS]
for t in range(T_out):
tok = toks[t] if t < len(toks) else EOS
logits[i, t, tok] = 100.0
return {"logits": logits, "halt_steps": None}
def _sieve35(n):
c = n + 1
while any(c % d == 0 for d in (2, 3, 5, 7) if d < c):
c += 1
return c
def _perfect(n):
return next(p for p in sieve_primes(500) if p > n)
def _easy_only(n):
return _perfect(n) if (n % 2 == 0 or n % 5 == 0) else 199
def test_probe_sieve35_classified_p1():
"""A pure {2,3,5,7} sieve must classify P1: errors are predictions of 121/143/169/187/209."""
r = probe_report(StubModel(_sieve35), DIGITS_CFG)
preds = sorted({e["pred"] for e in r["errors"]})
assert preds == sorted(FLAGGED_SIEVE_PREDS), r["errors"]
assert len(r["errors"]) == 22 # n in 113..120, 139..142, 167..168, 181..186, 199..200
assert r["code"] == "P1", r
def test_probe_perfect_classified_p3():
r = probe_report(StubModel(_perfect), DIGITS_CFG)
assert r["acc"] == 1.0 and r["code"] == "P3", r
def test_probe_easy_only_classified_p2():
r = probe_report(StubModel(_easy_only), DIGITS_CFG)
assert r["code"] == "P2", r
def test_probe_always_wrong_classified_p4():
r = probe_report(StubModel(lambda n: n), DIGITS_CFG)
assert r["code"] == "P4", r
def _write_csv(path, train_ems, val_ems):
with open(path, "w", newline="") as fh:
w = csv.writer(fh)
w.writerow(["step", "train_loss", "train_token_acc", "train_em", "val_token_acc",
"val_em", "mean_halt_steps", "log_examples", "param_count"])
for i, (t, v) in enumerate(zip(train_ems, val_ems)):
w.writerow([i * 200, 0.1, 1.0, t, 1.0, v, 10.0, "", 1000])
def _sig_of(tmp, tr, va):
p = os.path.join(tmp, "m.csv")
_write_csv(p, tr, va)
return grokking_signature(p)
def test_signature_o1_sharp_transition(tmp_path):
tr = [0.99] * 20
va = [0.1] * 15 + [0.2, 0.95, 0.99, 1.0, 1.0]
r = _sig_of(str(tmp_path), tr, va)
assert r["code"] == "O1", r
def test_signature_o2_memorization(tmp_path):
assert _sig_of(str(tmp_path), [1.0] * 20, [0.1] * 20)["code"] == "O2"
def test_signature_o3_gradual(tmp_path):
tr = [1.0] * 20
va = [0.3, 0.45, 0.6, 0.75, 0.9, 0.95] + [1.0] * 14
assert _sig_of(str(tmp_path), tr, va)["code"] == "O3"
def test_signature_o4_train_fails(tmp_path):
assert _sig_of(str(tmp_path), [0.6] * 20, [0.2] * 20)["code"] == "O4"
def test_signature_opartial(tmp_path):
assert _sig_of(str(tmp_path), [1.0] * 20, [0.6] * 20)["code"] == "O-PARTIAL"
def test_halting_h1_collapse_and_h2_pin():
from src.model_api import build_model
cfg = Config(model="rnn")
m = build_model(cfg)
with torch.no_grad():
m.halt_head.bias.fill_(50.0)
assert halting_report(m, cfg)["code"] == "H1"
m2 = build_model(cfg)
with torch.no_grad():
m2.halt_head.bias.fill_(-50.0)
assert halting_report(m2, cfg)["code"] == "H2"
|
