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
|
"""Tests for CUDA device handling, chunked evaluation, and scaling to 1000 and 100_000."""
import torch
from src.config import Config, parse_args
from src.data import build_examples, get_splits, next_prime, sieve_primes
from src.model_api import build_model
from src.train import evaluate, resolve_device
def test_device_resolution():
assert resolve_device("cpu") == torch.device("cpu")
auto_dev = resolve_device("auto")
expected = torch.device("cuda" if torch.cuda.is_available() else "cpu")
assert auto_dev == expected
def test_config_device_args():
cfg = parse_args(["rnn", "0", "--device", "cpu", "--eval_batch_size", "256", "--use_amp", "True", "--compile_model", "True"])
assert cfg.device == "cpu"
assert cfg.eval_batch_size == 256
assert cfg.use_amp is True
assert cfg.compile_model is True
def test_bisect_next_prime_large_range():
primes = sieve_primes(100_100)
# Check boundary primes
assert next_prime(2, primes) == 3
assert next_prime(97, primes) == 101
assert next_prime(997, primes) == 1009
assert next_prime(99_991, primes) == 100_003
def test_chunked_evaluate_matches_single_batch():
"""Verify that chunked evaluation with eval_batch_size=4 gives identical results to full batch."""
cfg = Config(model="rnn", eval_batch_size=4)
model = build_model(cfg)
train_in, val_in = get_splits(cfg)
val_ex = build_examples(val_in, cfg)
# Chunked
tok_chunked, em_chunked, per_chunked = evaluate(model, val_ex, cfg)
# Single full batch
cfg_full = Config(model="rnn", eval_batch_size=len(val_ex))
tok_full, em_full, per_full = evaluate(model, val_ex, cfg_full)
assert tok_chunked == tok_full
assert em_chunked == em_full
assert per_chunked == per_full
def test_range_extension_1000():
cfg = Config(range_start=2, range_end=1000, model="transformer")
train_in, val_in = get_splits(cfg)
assert len(train_in) + len(val_in) == 999
train_ex = build_examples(train_in, cfg)
assert len(train_ex) == len(train_in)
model = build_model(cfg)
tok, em, per = evaluate(model, build_examples(val_in[:20], cfg), cfg)
assert 0.0 <= em <= 1.0
def test_range_extension_100k_data_shapes():
"""Test data generation and model forward at 100k range without OOM."""
cfg = Config(range_start=2, range_end=100_000, model="rnn")
# Take a small slice of inputs from 100k range
sample_inputs = [2, 42, 999, 50_000, 99_991]
examples = build_examples(sample_inputs, cfg)
assert len(examples) == 5
model = build_model(cfg)
tok, em, per = evaluate(model, examples, cfg)
assert 0.0 <= em <= 1.0
|
