| """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.') |
|
|
