| """Windowed Master prediction: positive overlap weights, global phase0 output.""" |
| import torch |
| from .master import window_starts |
|
|
| @torch.inference_mode() |
| def predict_features(model,features,device): |
| tokens=len(features)//20 |
| if tokens==0:return torch.empty(0,3) |
| raw=features[:tokens*20].reshape(tokens,20,12).transpose(1,2) |
| total=torch.zeros(tokens,3);weight=torch.zeros(tokens) |
| for start in window_starts(tokens): |
| value=raw[start:start+640][None].to(device);n=value.shape[1] |
| with torch.autocast(device_type=device.type,dtype=torch.bfloat16,enabled=device.type=='cuda'): |
| pred=model(value,torch.tensor([n],device=device),torch.tensor([start],device=device))[0].float().cpu() |
| i=torch.arange(n);w=torch.minimum(torch.minimum(i+1,n-i),torch.full((n,),160)).float()/160 |
| total[start:start+n]+=pred*w[:,None];weight[start:start+n]+=w |
| if not (weight>0).all():raise ValueError('uncovered inference tokens') |
| total=total/weight[:,None] |
| windows=tokens//10 |
| return total[:windows*10].reshape(windows,10,3).mean(1) |
|
|
