File size: 3,107 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 | """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.')
|