diff options
Diffstat (limited to 'src/data.py')
| -rw-r--r-- | src/data.py | 28 |
1 files changed, 23 insertions, 5 deletions
diff --git a/src/data.py b/src/data.py index 35682d3..e6ea57b 100644 --- a/src/data.py +++ b/src/data.py @@ -63,15 +63,33 @@ def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list 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).""" + primes = sieve_primes(cfg.range_end + 100) + max_target = next_prime(cfg.range_end, primes) + if cfg.vocab_mode == "integers": + return 1, 2 # [value], [value, EOS] + 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) - T_in = max(len(x) for x in xs) - T_out = max(len(y) for y in ys) + in_max, out_max = _global_lengths(cfg) B = len(examples) - x = torch.full((B, T_in), cfg.pad_id, dtype=torch.long) - y = torch.full((B, T_out), cfg.pad_id, dtype=torch.long) + 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, : len(xi)] = torch.tensor(xi, dtype=torch.long) + 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) |