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
|
"""Tests for the Phase-3 flags: task_mode=is_prime, train_frac, lr_decay."""
import torch
from src.config import Config, parse_args
from src.data import build_examples, decode_tokens, get_splits
from src.eval import probe_report
from src.model_api import build_model
from tests.test_eval_classification import StubModel
def _isprime_sieve(n):
"""Primality via {2,3,5,7} with the d<c guard — EXACT on [2,100] (every composite <=100
has a prime factor <=10), wrong only on no-small-factor composites like 121."""
if n < 2:
return 0
return int(all(n % d != 0 for d in (2, 3, 5, 7) if d < n))
def test_is_prime_examples():
cfg = Config(task_mode="is_prime")
ex = build_examples([2, 3, 4, 97, 100], cfg)
got = [decode_tokens(y, cfg) for _, y in ex]
assert got == [1, 1, 0, 1, 0], got
def test_is_prime_sieve_stub_perfect_in_range():
"""The {2,3,5,7} sieve is exact on [2,100] -> evaluate() must see EM 1.0."""
from src.train import evaluate
cfg = Config(task_mode="is_prime")
train_in, val_in = get_splits(cfg)
m = StubModel(_isprime_sieve)
for split in (train_in, val_in):
tok, em, _ = evaluate(m, build_examples(split, cfg), cfg)
assert em == 1.0 and tok == 1.0
def test_is_prime_probe_classified_p1():
"""On [101,200] the sieve errs exactly on {121,143,169,187} (classified prime) -> P1."""
cfg = Config(task_mode="is_prime")
r = probe_report(StubModel(_isprime_sieve), cfg)
assert sorted(e["n"] for e in r["errors"]) == [121, 143, 169, 187], r["errors"]
assert r["code"] == "P1", r["code"]
def test_train_frac_subsamples_train_only():
cfg = Config(train_frac=0.5)
tr, va = get_splits(cfg)
assert len(va) == 30 # val untouched (locked size)
assert len(tr) == 34 # round(69*0.5) = 34 (banker's rounding)
assert not set(tr) & set(va)
# deterministic across calls
tr2, va2 = get_splits(Config(train_frac=0.5))
assert tr == tr2 and va == va2
# subsample is a subset of the full train split
tr_full, va_full = get_splits(Config(train_frac=1.0))
assert set(tr) <= set(tr_full)
def test_integers_mode_eval_does_not_crash_probe(tmp_path):
"""Regression: integers-mode eval crashed in the probe (out-of-vocab inputs >= 104).
Train a tiny integers model, then run the REAL eval main against it — must not crash,
and the probe must be N/A (Addendum 5: D2 is scored on in-range val EM only)."""
import json
import os
import subprocess
import sys
import torch
from src.data import build_examples, get_splits, make_batch
from src.model_api import build_model
cfg = Config(vocab_mode="integers")
m = build_model(cfg)
train_in, val_in = get_splits(cfg)
train_ex = build_examples(train_in, cfg)
opt = torch.optim.AdamW(m.parameters(), lr=1e-3)
for _ in range(10):
ex = train_ex[:16]
b = make_batch(ex, cfg)
out = m(b["x"], b["y_in"])
loss = torch.nn.functional.cross_entropy(
out["logits"].reshape(-1, cfg.vocab), b["y"].clamp(max=cfg.vocab - 1).reshape(-1), reduction="none")
mask = b["y_mask"].float().reshape(-1)
loss = (loss * mask).sum() / mask.sum()
opt.zero_grad(); loss.backward(); opt.step()
od = tmp_path / "rnn" / "seed0"
os.makedirs(od, exist_ok=True)
cfg.save(str(od / "config.json"))
torch.save(m.state_dict(), str(od / "last.pt"))
torch.save(m.state_dict(), str(od / "best.pt"))
import csv
with open(str(od / "metrics.csv"), "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"])
w.writerow([200, 0.1, 1.0, 1.0, 1.0, 0.5, 3.0, "", 1000])
r = subprocess.run([sys.executable, "-m", "src.eval", "rnn", "0", "--runs-dir", str(tmp_path)],
capture_output=True, text=True, cwd=os.path.join(os.path.dirname(__file__), ".."))
assert r.returncode == 0, r.stderr[-600:]
res = json.load(open(str(od / "results.json")))
assert res["final"]["probe"]["code"] == "N/A"
def test_new_flags_parse():
a = parse_args(["rnn", "0", "--task_mode", "is_prime", "--train_frac", "0.4", "--lr_decay", "True"])
assert a.task_mode == "is_prime" and a.train_frac == 0.4 and a.lr_decay is True
b = parse_args(["transformer", "0", "--lr_decay", "False"])
assert b.lr_decay is False and b.task_mode == "next_prime" and b.train_frac == 1.0
|
