File size: 3,453 Bytes
317d342 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """Development-only frozen parent with learned service/log/base head selection."""
import torch
from torch import nn
from ._window_gate import WindowGateModel
class SourceKindGateModel(WindowGateModel):
"""Same parent forward/packing; two CLS gates choose one head per window."""
def __init__(self, encoder):
super().__init__(encoder)
self.log_head = nn.Linear(encoder.config.hidden_size, 1, bias=False)
self.log_gate = nn.Linear(encoder.config.hidden_size, 1)
def choose_head(self, pooled, cls):
gate = self.window_gate(cls)
original = self.head(pooled)
service = self.service_head(pooled)
log_gate = self.log_gate(cls)
log = self.log_head(pooled) + self.head.bias
valid = (torch.isfinite(gate) & torch.isfinite(log_gate) & torch.isfinite(original)
& torch.isfinite(service) & torch.isfinite(log))
chosen = torch.where(gate >= 0, service, torch.where(log_gate >= 0, log, original))
return torch.where(valid, chosen, torch.full_like(chosen, float('nan')))
def self_check():
"""Synthetic CPU branch/pooling/state regression; no trained encoder or CUDA."""
from types import SimpleNamespace
class Encoder(nn.Module):
config = SimpleNamespace(hidden_size=1)
def forward(self, input_ids, **kwargs):
return SimpleNamespace(last_hidden_state=input_ids.float().unsqueeze(-1))
model = SourceKindGateModel(Encoder()).eval()
with torch.no_grad():
model.head.weight.fill_(10); model.head.bias.fill_(2)
model.service_head.weight.fill_(-10); model.service_head.bias.fill_(3)
model.window_gate.weight.fill_(1); model.window_gate.bias.zero_()
model.log_head.weight.fill_(4)
model.log_gate.weight.fill_(-1); model.log_gate.bias.fill_(-1)
for first, expected in [(-2., 4.), (-1., 4.), (-.5, 7.), (0., -2.), (2., -2.)]:
inputs = torch.tensor([[first, -1., 2.]])
result = model(inputs, torch.ones_like(inputs), [(1, 3)])
assert result.tolist() == [expected]
states = inputs.unsqueeze(-1)
tokens = model.choose_head(states, states[:, :1, :]).squeeze(-1)
assert torch.equal(tokens[:, 1:3].mean(1), result)
state = {k: v.detach().clone() for k, v in model.state_dict().items()}
for key in state:
model.load_state_dict(state)
with torch.no_grad():
dict(model.named_parameters())[key].fill_(float('inf'))
try:
model(torch.tensor([[-.5, -1., 2.]]), torch.ones(1, 3), [(1, 3)])
except ValueError as error:
assert str(error) == 'nonfinite_scores'
else:
raise AssertionError('even_unselected_nonfinite_branches_must_refuse')
model.load_state_dict(state)
with torch.no_grad():
model.log_gate.bias.fill_(100)
assert model(torch.tensor([[2., -1., 2.]]), torch.ones(1, 3), [(1, 3)]).tolist() == [-2.]
parent = WindowGateModel(Encoder())
extra = {'log_head.weight', 'log_gate.weight', 'log_gate.bias'}
assert set(state) - set(parent.state_dict()) == extra
encoder = Encoder(); encoder.config = SimpleNamespace(hidden_size=768)
assert sum(v.numel() for k, v in SourceKindGateModel(encoder).state_dict().items() if k in extra) == 1537
if __name__ == '__main__':
self_check()
print('Source-kind boundaries, service priority, pooling, extra state and nonfinite refusal pass.')
|