File size: 3,000 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Shared learned token composition, with no numeric parser in inference."""
import re
import torch
from torch import Tensor, nn
MAX_NUMBER_TOKENS = 8
def number_token_ids(tokenizer, ids: tuple[int, ...]) -> tuple[int, ...]:
"""Route only the raw Steps text; never convert it to a number.
One token per character is a tokenizer contract, not a digit-value lookup.
The public prompt and executor routing remain unchanged.
"""
text = tokenizer.decode(list(ids), skip_special_tokens=False,
clean_up_tokenization_spaces=False)
match = re.fullmatch(r'Rules: [^\n]+\nStart: [A-Z]\nSteps: ([1-9][0-9]*)\nAnswer:', text)
if match is None:
raise ValueError('Number reader requires the raw Rules/Start/Steps/Answer prompt')
field = match.group(1)
tokens = tokenizer.encode(field, add_special_tokens=False)
if not 1 <= len(tokens) <= MAX_NUMBER_TOKENS or len(tokens) != len(field):
raise ValueError('Number reader requires one token per decimal character and at most eight characters')
# Length alone does not rule out a tokenizer that groups/reorders characters.
if any(tokenizer.decode([t], clean_up_tokenization_spaces=False) != c
for t, c in zip(tokens, field, strict=True)):
raise ValueError('Number reader tokenizer does not preserve individual decimal characters')
return tuple(tokens)
class SharedNumberReader(nn.Module):
"""Read padded frozen embeddings [B,8*(H+1)] using the same map at each slot.
Last channel of each slot is a validity mask. Valid tokens are left aligned.
The gain starts at identity, not a decimal radix. Training labels fit both
gain and token projection. Accumulate in float64 to reduce numerical error,
then return the controller's scalar-memory dtype [B,1].
"""
def __init__(self, width: int) -> None:
super().__init__()
self.width = width
self.in_features = MAX_NUMBER_TOKENS * (width + 1)
self.token_value = nn.Linear(width, 1)
self.gain = nn.Parameter(torch.ones(()))
def unpack(self, context: Tensor) -> tuple[Tensor, Tensor]:
if context.ndim != 2 or context.shape[1] != self.in_features:
raise ValueError('Need padded number embeddings with token-validity masks')
slots = context.reshape(-1, MAX_NUMBER_TOKENS, self.width + 1)
return slots[..., :self.width], slots[..., self.width:]
def forward(self, context: Tensor) -> Tensor:
embeddings, mask = self.unpack(context.detach())
values = torch.nn.functional.linear(embeddings.double(), self.token_value.weight.double(),
self.token_value.bias.double())
state = values.new_zeros(len(values), 1)
for value, valid in zip(values.unbind(1), mask.unbind(1), strict=True):
state = torch.where(valid.bool(), self.gain.double() * state + value, state)
return state.to(self.token_value.weight.dtype)
|