File size: 1,079 Bytes
e809d15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | """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)
|