summaryrefslogtreecommitdiff
path: root/tests/test_eval_classification.py
blob: feffe60987f82e30e62be5c953bc4f4499c4543b (plain)
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
136
137
138
139
"""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 _sieve_rank_signature, 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_p5():
    """A pure {2,3,5,7} sieve must classify P5(4): errors match the rank-4 sieve signature."""
    r = probe_report(StubModel(_sieve35), DIGITS_CFG)
    preds = sorted({e["pred"] for e in r["errors"]})
    # all error preds must be composites with all factors > 7 (rank-4 signature)
    # use extended range to cover predictions that land outside [101,200]
    sig4_extended = _sieve_rank_signature(4, 101, 250)
    assert all(p in sig4_extended for p in preds), f"unexpected preds: {preds}"
    assert len(r["errors"]) == 22
    assert r["code"] == "P5(4)", r
    assert r["sieve_rank"] == 4


def test_probe_perfect_classified_p6():
    r = probe_report(StubModel(_perfect), DIGITS_CFG)
    assert r["acc"] == 1.0 and r["code"] == "P6", 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"