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
|
"""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_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
|
