File size: 1,102 Bytes
e29f88c | 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 | """Hidden-state-only completion head for explicit-step pointer recurrence."""
import torch
from torch import Tensor, nn
class CompletionHead(nn.Module):
"""One stop logit per answer-position recurrent state [B,H].
No numeric loop index, parsed requested depth, decoded symbol, or target is
accepted by this interface. The prompt can encode its requested Steps.
"""
def __init__(self, width: int, intermediate: int = 128) -> None:
super().__init__()
if width < 1 or intermediate < 1:
raise ValueError("Completion-head widths must be positive")
self.width = width
self.intermediate = intermediate
self.network = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate),
nn.GELU(), nn.Linear(intermediate, 1))
def forward(self, answer_state: Tensor) -> Tensor:
if answer_state.ndim != 2 or answer_state.shape[-1] != self.width:
raise ValueError("Completion head requires answer states [batch, hidden]")
return self.network(answer_state).squeeze(-1)
|