Spaces:
Sleeping
Sleeping
| """MS MARCO MiniLM cross-encoder, loaded from uploaded weights. CPU, offline.""" | |
| import os,unicodedata,re,time | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from safetensors.torch import load_file | |
| from minilm_cpu_pipeline import BertEmb,Enc,WordPiece | |
| from neural_evidence_graph_core import ROOT | |
| class BertCrossEncoder(nn.Module): | |
| def __init__(self): | |
| super().__init__();self.embeddings=BertEmb();self.encoder=Enc() | |
| self.pooler=nn.Linear(384,384);self.classifier=nn.Linear(384,1) | |
| def forward(self,ids,mask,types): | |
| pos=self.embeddings.position_ids[:,:ids.shape[1]] | |
| h=self.embeddings.LayerNorm(self.embeddings.word_embeddings(ids)+self.embeddings.position_embeddings(pos)+self.embeddings.token_type_embeddings(types)) | |
| h=self.encoder(h,mask[:,None,None,:].bool()) | |
| return self.classifier(torch.tanh(self.pooler(h[:,0,:]))).squeeze(-1) | |
| def load_crossencoder(): | |
| torch.set_num_threads(min(4,os.cpu_count() or 4)) | |
| try:torch.set_num_interop_threads(1) | |
| except RuntimeError:pass | |
| m=BertCrossEncoder();raw=load_file(str(ROOT/'model.safetensors'),device='cpu');clean={} | |
| for k,v in raw.items(): | |
| if k=='bert.embeddings.position_ids':continue | |
| if k.startswith('bert.pooler.dense.'):k='pooler.'+k[len('bert.pooler.dense.'):] | |
| elif k.startswith('bert.'):k=k[5:] | |
| clean[k]=v | |
| m.load_state_dict(clean,strict=True);m.eval() | |
| print('Loaded original pretrained MS MARCO model, tensors',len(clean),flush=True) | |
| return m | |
| class PairTokenizer(WordPiece): | |
| def __init__(self,length=256):super().__init__(str(ROOT/'vocab.txt'),length) | |
| def tokens(self,text): | |
| text=''.join(ch for ch in unicodedata.normalize('NFD',text.lower()) if unicodedata.category(ch)!='Mn' and (ch.isspace() or unicodedata.category(ch)[0]!='C')) | |
| out=[] | |
| for x in re.findall(r'[^\W_]+|_|[^\w\s]',text,flags=re.UNICODE):out.extend(self.wp(x)) | |
| return out | |
| def pair(self,q,d): | |
| a=self.tokens(q);b=self.tokens(d);cap=self.max_length-3 | |
| while len(a)+len(b)>cap: | |
| if len(b)>len(a):b.pop() | |
| else:a.pop() | |
| ids=[101]+a+[102]+b+[102] | |
| tt=[0]*(len(a)+2)+[1]*(len(b)+1) | |
| return ids,tt | |
| def score_pairs(m,tok,pairs,batch=20): | |
| out=np.empty(len(pairs),np.float32);times=[] | |
| for start in range(0,len(pairs),batch): | |
| p=pairs[start:start+batch];inputs=[tok.pair(q,d) for q,d in p] | |
| L=max(len(x[0]) for x in inputs) | |
| ids=np.zeros((len(p),L),np.int64);tt=np.zeros_like(ids);mask=np.zeros_like(ids) | |
| for i,(u,v) in enumerate(inputs):ids[i,:len(u)]=u;tt[i,:len(v)]=v;mask[i,:len(u)]=1 | |
| t=time.perf_counter();out[start:start+len(p)]=m(torch.from_numpy(ids),torch.from_numpy(mask),torch.from_numpy(tt)).numpy();times.append(time.perf_counter()-t) | |
| if start and start%(batch*35)==0:print('scored',start,'/',len(pairs),flush=True) | |
| return out,{'total_forward_s':float(sum(times)),'forward_ms_per_pair':float(sum(times)*1000/max(1,len(pairs))),'batches':len(times)} | |