diff options
Diffstat (limited to 'src/model_api.py')
| -rw-r--r-- | src/model_api.py | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/src/model_api.py b/src/model_api.py index 07d5d8a..c82baf4 100644 --- a/src/model_api.py +++ b/src/model_api.py @@ -41,7 +41,15 @@ def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int x = pad_inputs(x, cfg) from src.data import _global_lengths _, out_max = _global_lengths(cfg) - max_len = max_len or max(getattr(cfg, "max_out_len", 6), out_max) + if max_len is None: + max_len = max(getattr(cfg, "max_out_len", 6), out_max) + if cfg.scratch_mode == "structured": + # Probe-window traces ([range_end+1, +1000]) exceed the in-range layout + # bound (measured: 633-token trace at n=1327 vs out_max=369). Cap at the + # model's positional table instead so OOD traces decode fully. + pos = getattr(model, "pos", None) + if pos is not None: + max_len = max(max_len, pos.num_embeddings - x.shape[1] - 1) B = x.shape[0] y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device) for _ in range(max_len): |