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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
"""Prime dataset: n -> next prime, digit-tokenized; splits and batching."""
import bisect
import random
import torch
from src.config import Config
def sieve_primes(limit: int) -> list[int]:
"""All primes <= limit (inclusive)."""
if limit < 2:
return []
is_prime = [True] * (limit + 1)
is_prime[0] = is_prime[1] = False
for p in range(2, int(limit ** 0.5) + 1):
if is_prime[p]:
for m in range(p * p, limit + 1, p):
is_prime[m] = False
return [i for i in range(2, limit + 1) if is_prime[i]]
def next_prime(n: int, primes: list[int]) -> int:
idx = bisect.bisect_right(primes, n)
if idx < len(primes):
return primes[idx]
raise ValueError(f"no prime > {n} in supplied list")
def encode_int(n: int, cfg: Config) -> list[int]:
if cfg.vocab_mode == "integers":
return [n]
return [int(d) for d in str(n)]
def decode_tokens(ts, cfg: Config) -> int:
"""Decode a token sequence, stopping at EOS. -1 if nothing decodable."""
if cfg.vocab_mode == "integers":
return int(ts[0]) if len(ts) else -1
digits = []
for t in ts:
if t == cfg.eos_id:
break
digits.append(int(t))
return int("".join(map(str, digits))) if digits else -1
def is_prime_n(n: int) -> bool:
"""Exact primality for n >= 2."""
if n < 2:
return False
d = 2
while d * d <= n:
if n % d == 0:
return False
d += 1
return True
def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
"""(train, val) input lists, seeded shuffle, no overlap. train_frac subsamples the
TRAIN split only (E5); the val split is untouched (its size is locked by prereg)."""
rng = random.Random(cfg.seed)
inputs = list(range(cfg.range_start, cfg.range_end + 1))
rng.shuffle(inputs)
n_val = max(1, round(len(inputs) * cfg.holdout_frac))
train, val = sorted(inputs[n_val:]), sorted(inputs[:n_val])
if cfg.train_frac < 1.0:
n_tr = max(1, round(len(train) * cfg.train_frac))
# deterministic subsample: seeded shuffle, take first n_tr
sub = random.Random(cfg.seed + 1000) # distinct stream from split shuffle
sub.shuffle(train)
train = sorted(train[:n_tr])
return train, val
def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]:
"""[(input_tokens, target_tokens+EOS), ...]. task_mode selects the target function."""
margin = max(100, int(cfg.range_end * 0.05) + 50)
primes = sieve_primes(cfg.range_end + margin) if cfg.task_mode != "is_prime" else []
out = []
for n in inputs:
if cfg.task_mode == "is_prime":
p = 1 if is_prime_n(n) else 0
else:
p = next_prime(n, primes)
out.append((encode_int(n, cfg), encode_int(p, cfg) + [cfg.eos_id]))
return out
def _global_lengths(cfg: Config) -> tuple[int, int]:
"""(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix)."""
if cfg.vocab_mode == "integers":
return 1, 2 # [value], [value, EOS]
if cfg.task_mode == "is_prime":
return len(str(cfg.range_end)), 2 # digit token "1"/"0" + EOS
margin = max(100, int(cfg.range_end * 0.05) + 50)
primes = sieve_primes(cfg.range_end + margin)
max_target = next_prime(cfg.range_end, primes)
return len(str(cfg.range_end)), len(str(max_target)) + 1 # digits + EOS
def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor:
"""LEFT-pad inputs to the global in_max so absolute positions are layout-invariant."""
in_max, _ = _global_lengths(cfg)
if x.shape[1] < in_max:
pad = torch.full((x.shape[0], in_max - x.shape[1]), cfg.pad_id, dtype=x.dtype, device=x.device)
x = torch.cat([pad, x], dim=1)
return x
def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]:
"""Fixed global layout (not batch-max): x left-padded to in_max, y right-padded to out_max."""
xs, ys = zip(*examples)
in_max, out_max = _global_lengths(cfg)
B = len(examples)
x = torch.full((B, in_max), cfg.pad_id, dtype=torch.long)
y = torch.full((B, out_max), cfg.pad_id, dtype=torch.long)
for i, (xi, yi) in enumerate(examples):
x[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long)
y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long)
# teacher-forced decoder input: BOS(=EOS reuse) then shifted y
y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1)
y_mask = y != cfg.pad_id
return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask}
|
