| """Prototype: frozen-parent encoder/base head with a learned CLS-window selector.""" |
| import torch |
| from torch import nn |
| from tokenfold_select._model import JointGroupModel |
|
|
|
|
| class WindowGateModel(JointGroupModel): |
| """One head choice per complete encoder window, not per token or group.""" |
| def __init__(self, encoder): |
| super().__init__(encoder) |
| self.window_gate=nn.Linear(encoder.config.hidden_size,1) |
| self.service_head=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) |
| valid=torch.isfinite(gate)&torch.isfinite(original)&torch.isfinite(service) |
| chosen=torch.where(gate>=0,service,original) |
| return torch.where(valid,chosen,torch.full_like(chosen,float('nan'))) |
|
|
| def forward(self, input_ids, attention_mask, ranges): |
| if input_ids.ndim!=2 or input_ids.shape[0]!=1 or attention_mask.shape!=input_ids.shape: |
| raise ValueError('single_sequence_required') |
| if not 1<=len(ranges)<=512: |
| raise ValueError('invalid_ranges') |
| previous=0 |
| for start,end in ranges: |
| if not previous<=start<end<=input_ids.shape[1] or not attention_mask[0,start:end].all(): |
| raise ValueError('invalid_ranges') |
| previous=end |
| positions=torch.arange(input_ids.shape[1],device=input_ids.device).unsqueeze(0) |
| states=self.encoder(input_ids=input_ids,attention_mask=attention_mask,position_ids=positions).last_hidden_state[0] |
| pooled=torch.stack([states[a:b].mean(0) for a,b in ranges]) |
| scores=self.choose_head(pooled,states[:1]).squeeze(-1) |
| if not torch.isfinite(scores).all(): |
| raise ValueError('nonfinite_scores') |
| return scores |
|
|
|
|
| def self_check(): |
| 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)) |
| m=WindowGateModel(Encoder()).eval() |
| with torch.no_grad(): |
| m.head.weight.fill_(10);m.head.bias.zero_() |
| m.service_head.weight.fill_(-10);m.service_head.bias.zero_() |
| m.window_gate.weight.fill_(1);m.window_gate.bias.zero_() |
| for first,expected in [(-2.,5.),(2.,-5.),(0.,-5.)]: |
| x=torch.tensor([[first,-1.,2.]]) |
| result=m(x,torch.ones_like(x),[(1,3)]) |
| assert result.tolist()==[expected] |
| states=x.unsqueeze(-1) |
| tokens=m.choose_head(states,states[:,:1,:]).squeeze(-1) |
| assert torch.equal(tokens[:,1:3].mean(1),result) |
| with torch.no_grad():m.service_head.weight.fill_(float('inf')) |
| try:m(torch.tensor([[-2.,-1.,2.]]),torch.ones(1,3),[(1,3)]) |
| except ValueError as error:assert str(error)=='nonfinite_scores' |
| else:raise AssertionError('masked nonfinite branch must refuse, not fall back') |
|
|
|
|
| if __name__=='__main__':self_check();print('CLS-global head choice, pooling parity, zero-boundary and nonfinite refusal checks passed.') |
|
|
