diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/test_data.py | 78 | ||||
| -rw-r--r-- | tests/test_models.py | 77 |
2 files changed, 155 insertions, 0 deletions
diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..fe84e94 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,78 @@ +import torch + +from src.config import Config +from src.data import (build_examples, decode_tokens, encode_int, get_splits, + make_batch, next_prime, sieve_primes) + + +def test_sieve_primes_known(): + assert sieve_primes(30) == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29] + assert sieve_primes(1) == [] + assert sieve_primes(2) == [2] + + +def test_next_prime(): + p = sieve_primes(200) + assert next_prime(42, p) == 43 + assert next_prime(100, p) == 101 + assert next_prime(2, p) == 3 + assert next_prime(97, p) == 101 + + +def test_encode_decode_roundtrip_digits(): + cfg = Config(vocab_mode="digits") + for n in [2, 7, 10, 42, 99, 100]: + assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n + assert encode_int(42, cfg) == [4, 2] + + +def test_encode_decode_roundtrip_integers(): + cfg = Config(vocab_mode="integers", range_end=100) + for n in [2, 42, 100]: + assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n + assert cfg.vocab == 102 # 0..101 + EOS + + +def test_splits_sizes_and_no_overlap(): + cfg = Config() + tr, va = get_splits(cfg) + assert len(tr) == 69 and len(va) == 30 + assert not set(tr) & set(va) + assert set(tr) | set(va) == set(range(2, 101)) + + +def test_splits_seed_stable(): + a1, b1 = get_splits(Config(seed=3)) + a2, b2 = get_splits(Config(seed=3)) + assert a1 == a2 and b1 == b2 + a3, _ = get_splits(Config(seed=4)) + assert a1 != a3 + + +def test_build_examples_targets(): + cfg = Config() + ex = build_examples([2, 42, 99, 100], cfg) + assert [decode_tokens(y, cfg) for _, y in ex] == [3, 43, 101, 101] + assert ex[0][1][-1] == cfg.eos_id # EOS-terminated + + +def test_make_batch_shapes_and_padding(): + cfg = Config() + ex = build_examples([2, 7, 42, 99], cfg) + b = make_batch(ex, cfg) + assert b["x"].shape == (4, 2) # max input len 2 (42, 99) + assert b["y"].shape == (4, 4) # max output len 4 ("101"+EOS) + assert b["y_in"].shape == (4, 4) + assert b["y_mask"].shape == (4, 4) + # pad positions masked out + assert not b["y_mask"][0, 3] # "3"+EOS row has pad at index 3 + assert b["y_mask"][0, 1] + # BOS position of y_in is eos_id + assert (b["y_in"][:, 0] == cfg.eos_id).all() + + +def test_flagged_composites_are_composite(): + # sanity: the diagnostic set in eval.py really is composite and needs divisors > 7 + for n in (121, 143, 169, 187): + assert any(n % d == 0 for d in range(2, int(n ** 0.5) + 1)) + assert all(n % d != 0 for d in (2, 3, 5, 7)) diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..7994e86 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,77 @@ +import torch +import torch.nn.functional as F + +from src.config import Config +from src.data import build_examples, make_batch +from src.model_api import build_model + + +def _batch(cfg, inputs=(2, 7, 42, 99)): + ex = build_examples(list(inputs), cfg) + return make_batch(ex, cfg) + + +def test_rnn_shapes_and_halting_range(): + cfg = Config(model="rnn", max_steps=20, min_steps=2) + m = build_model(cfg) + b = _batch(cfg) + out = m(b["x"], b["y_in"]) + assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab) + s = out["halt_steps"] + assert s.shape == (4,) + assert bool((s >= cfg.min_steps).all()) + assert bool((s <= cfg.max_steps).all()) + + +def test_rnn_fixed_k_ablation(): + cfg = Config(model="rnn", halting=False, max_steps=20) + m = build_model(cfg) + b = _batch(cfg) + out = m(b["x"], b["y_in"]) + assert torch.allclose(out["halt_steps"], torch.full((4,), 20.0)) + + +def test_transformer_shapes(): + cfg = Config(model="transformer") + m = build_model(cfg) + b = _batch(cfg) + out = m(b["x"], b["y_in"]) + assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab) + assert out["halt_steps"] is None + + +def test_transformer_integers_mode(): + cfg = Config(model="transformer", vocab_mode="integers") + m = build_model(cfg) + b = _batch(cfg) + out = m(b["x"], b["y_in"]) + assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab) + + +def test_param_counts_logged_not_gated(): + r = build_model(Config(model="rnn")).param_count() + t = build_model(Config(model="transformer")).param_count() + assert r > 1000 and t > 1000 + + +def test_forward_backward_no_nan_both_models(): + for model_name in ("rnn", "transformer"): + cfg = Config(model=model_name) + m = build_model(cfg) + b = _batch(cfg) + out = m(b["x"], b["y_in"]) + y_safe = b["y"].clamp(max=cfg.vocab - 1) + loss = F.cross_entropy(out["logits"].reshape(-1, cfg.vocab), y_safe.reshape(-1), reduction="none") + mask = b["y_mask"].float().reshape(-1) + loss = (loss * mask).sum() / mask.sum() + loss.backward() + assert torch.isfinite(loss).item(), f"{model_name} loss NaN" + + +def test_greedy_decode_shape(): + from src.model_api import greedy_decode + cfg = Config(model="rnn") + m = build_model(cfg) + x = torch.tensor([[4, 2], [9, 9]], dtype=torch.long) + gen = greedy_decode(m, x, cfg) + assert gen.shape == (2, cfg.max_out_len) |