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
|
"""Model interface contract + build dispatch + shared greedy decode."""
import torch
import torch.nn as nn
from src.config import Config
from src.data import pad_inputs
class PrimeModel(nn.Module):
"""Contract: forward(x, y_in) -> {"logits": (B,T_out,vocab), "halt_steps": (B,) or None}."""
def forward(self, x, y_in):
raise NotImplementedError
def param_count(self) -> int:
return sum(p.numel() for p in self.parameters())
def build_model(cfg: Config) -> PrimeModel:
if cfg.model == "rnn":
from src.models.rnn import TiedRNN
return TiedRNN(cfg)
if cfg.model == "transformer":
from src.models.transformer import TransformerBaseline
return TransformerBaseline(cfg)
raise ValueError(f"unknown model: {cfg.model}")
@torch.no_grad()
def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int | None = None) -> torch.Tensor:
"""Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped).
Inputs are LEFT-padded to the global layout so positions match training exactly
(batch/singleton invariance — codex BLOCKER fix)."""
x = pad_inputs(x, cfg)
max_len = max_len or cfg.max_out_len
B = x.shape[0]
y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device)
for _ in range(max_len):
out = model(x, y_in)
nxt = out["logits"][:, -1].argmax(-1)
y_in = torch.cat([y_in, nxt[:, None]], dim=1)
return y_in[:, 1:]
|
