spb/wpai0001_uqo

1"""Greedy scratchpad decoding for toy models (batched, per-row prompt lengths)."""2from __future__ import annotations34import torch56from wp_ai_0001.models.toy_transformer import ToyTransformer789@torch.no_grad()10def greedy_decode(model: ToyTransformer, ids: torch.Tensor, prompt_len: torch.Tensor, n_answer: torch.Tensor, K: int) -> torch.Tensor:11    """Fill answer positions greedily; returns predicted answer ids [B,K] (pad beyond n_answer).1213    ``ids`` must contain the prompt; answer positions are overwritten step by step so the14    generated prefix conditions later steps (no teacher forcing).15    """16    ids = ids.clone()17    B = ids.shape[0]18    ar = torch.arange(B, device=ids.device)19    out = torch.zeros(B, K, dtype=torch.long, device=ids.device)20    for s in range(K):21        read = prompt_len - 1 + s22        logits = model(ids, read)  # hook positions = the read position for this step23        pred = logits[ar, read].argmax(-1)24        write = prompt_len + s25        active = s < n_answer26        ids[ar[active], write[active]] = pred[active]27        out[:, s] = torch.where(active, pred, torch.zeros_like(pred))28    return out29