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
|
"""Batched evaluate() correctness — the speedup rewrite must preserve exact semantics."""
import torch
from src.config import Config
from src.data import build_examples, get_splits
from src.train import evaluate
from tests.test_eval_classification import StubModel, _sieve35
def test_batched_evaluate_known_model():
"""The {2,3,5,7} sieve is CORRECT on the whole [2,100] range (any composite <= 101 has a
factor in {2,3,5,7}, and every in-range next prime has no small divisor) — so evaluate()
against the sieve stub must return token acc 1.0 and exact-match 1.0 on both splits."""
cfg = Config()
train_in, val_in = get_splits(cfg)
m = StubModel(_sieve35)
for split_in in (train_in, val_in):
ex = build_examples(split_in, cfg)
tok, em, per = evaluate(m, ex, cfg)
assert tok == 1.0, f"token acc {tok} (split size {len(ex)})"
assert em == 1.0, f"exact-match {em} (split size {len(ex)})"
assert len(per) == len(ex)
assert all(ok for _, _, _, ok in per)
def test_batched_evaluate_matches_singleton_greedy():
"""For a REAL model, batched evaluate's greedy decode must equal greedy_decode per row
(batch-invariance guarantees this — belt and braces against a regressed layout fix)."""
from src.model_api import build_model, greedy_decode
from src.data import make_batch, decode_tokens
cfg = Config(model="rnn")
torch.manual_seed(0)
m = build_model(cfg)
_, val_in = get_splits(cfg)
val_ex = build_examples(val_in, cfg)
_, em, per = evaluate(m, val_ex, cfg)
# independent per-example greedy via the public helper
for i, (x, target) in enumerate(val_ex):
xi = torch.tensor(x, dtype=torch.long).unsqueeze(0)
gen = greedy_decode(m, xi, cfg)[0].tolist()
pred_alone = decode_tokens(gen, cfg)
assert per[i][2] == pred_alone, f"example {i}: batched {per[i][2]} vs singleton {pred_alone}"
|