File size: 2,462 Bytes
a82e638 | 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 | 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()
|