| import hashlib |
| import json |
| import random |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from .interaction_head import TriVectorInteractionHead |
| from .verifier_batching import prepare_batch_inputs |
|
|
| def identity(value): |
| return hashlib.sha256(json.dumps(value,sort_keys=True,separators=(',',':'),allow_nan=False).encode()).hexdigest() |
|
|
| def seed_all(seed): |
| random.seed(seed);np.random.seed(seed);torch.manual_seed(seed) |
| if torch.cuda.is_available():torch.cuda.manual_seed_all(seed) |
| torch.set_num_threads(8);torch.use_deterministic_algorithms(True) |
| torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False |
|
|
| class FullModel(nn.Module): |
| def __init__(self,backbone): |
| super().__init__();self.backbone=backbone |
| self.head=TriVectorInteractionHead(backbone.config.hidden_size).to(next(backbone.parameters()).device) |
| def forward(self,input_ids,attention_mask,position_ids,segments): |
| hidden=self.backbone(input_ids=input_ids,attention_mask=attention_mask, |
| position_ids=position_ids,use_cache=False,return_dict=True).last_hidden_state |
| pools=self.head.pool_segments(hidden,segments,attention_mask) |
| return self.head(*(F.layer_norm(p.float(),(p.shape[-1],)) for p in pools)) |
|
|
| def collate(rows,pad,device): |
| ids,mask=prepare_batch_inputs([torch.tensor(r['input_ids'],device=device) for r in rows],pad) |
| seg=torch.zeros_like(ids) |
| for i,row in enumerate(rows):seg[i,-len(row['segments']):]=torch.tensor(row['segments'],device=device) |
| return dict(input_ids=ids,attention_mask=mask,position_ids=(mask.cumsum(-1)-1).clamp_min(0),segments=seg) |
|
|
| def infer(model,rows,pad,device='cuda',progress=None): |
| model.eval();scores=[] |
| with torch.inference_mode(): |
| for i,row in enumerate(rows): |
| value=model(**collate([row],pad,device))[0].float().cpu() |
| if not torch.isfinite(value).all():raise RuntimeError('Nonfinite inference') |
| scores.append(value) |
| if progress and (i%25==0 or i+1==len(rows)):progress(i+1,len(rows)) |
| return torch.stack(scores) |
|
|
| def tensor_hash(state): |
| h=hashlib.sha256() |
| for name,p in sorted(state.items()): |
| t=p.detach().cpu().contiguous() |
| h.update(name.encode());h.update(str(t.dtype).encode());h.update(str(tuple(t.shape)).encode()) |
| h.update(t.reshape(-1).view(torch.uint8).numpy().tobytes()) |
| return h.hexdigest() |
|
|
