Spaces:
Sleeping
Sleeping
| """Pretrained MS MARCO MiniLM-L6 cross-encoder CPU-only held-out evaluation. | |
| Inputs are user-uploaded safetensors/tokenizer and existing HotpotQA / 2Wiki candidate pools. | |
| A fixed 200/600 split (seed 20261030) is reused from previous experiments. | |
| No gold labels used for the cross-encoder, and none fed to selection. | |
| """ | |
| import os,time,json,argparse,statistics,unicodedata,re | |
| os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='4' | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from safetensors.torch import load_file | |
| from minilm_cpu_pipeline import BertEmb,Enc,WordPiece | |
| from hotpot_subset_eval import load_data as hotpot_data | |
| from wiki_subset_eval import load_data as wiki_data | |
| from graph_bridge_eval import prep_links,graph_select | |
| ROOT=Path('/mnt/data');RNG=20261030 | |
| 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,token_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(token_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)); torch.set_num_interop_threads(1) | |
| 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 | |
| missing,unexpected=m.load_state_dict(clean,strict=True) | |
| assert not missing and not unexpected,(missing,unexpected) | |
| m.eval();print('Loaded pretrained MS MARCO cross-encoder: 106 tensors, all parameters match',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=16,limit=None): | |
| 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(f' scored {start}/{len(pairs)} pairs',flush=True) | |
| return out,{'total_forward_s':float(sum(times)),'forward_ms_per_pair':float(sum(times)*1000/max(len(pairs),1)),'batches':len(times)} | |
| def coverage(scores,items,k): | |
| y=np.zeros((len(items),10),np.uint8) | |
| for i,r in enumerate(items): | |
| gold={x.casefold().strip() for x,_ in r['supporting_facts']} | |
| for j,(title,_) in enumerate(r['context']):y[i,j]=int(title.casefold().strip() in gold) | |
| orders=np.argsort(-scores,axis=1,kind='stable') | |
| return np.asarray([int(y[i,orders[i,:k]].sum()==y[i].sum() and y[i].sum()>0) for i in range(len(items))]),orders | |
| def run(ds,m,tok,batch=16,max_q=None): | |
| start=time.perf_counter() | |
| items,_=(hotpot_data() if ds=='hotpot' else wiki_data()) | |
| test=items[200:800];n=len(test) | |
| if max_q is not None:n=min(n,max_q);test=test[:n] | |
| pairs=[(r['question'],title+'. '+ ' '.join(sentences)) for r in test for title,sentences in r['context']] | |
| print(f'{ds}: scoring {n} heldout questions, {len(pairs)} pairs, seq_limit={tok.max_length}',flush=True) | |
| ts=time.perf_counter();ranks,perf=score_pairs(m,tok,pairs,batch=batch) | |
| ts=time.perf_counter()-ts; scores=ranks.reshape(n,10) | |
| emb=np.load(ROOT/('hotpot_subset_out/hotpot_minilm_800.npz' if ds=='hotpot' else 'wiki_subset_out/wiki_minilm_800.npz')) | |
| Q=emb['q'][200:200+n];D=emb['docs'][200:200+n];Fq=emb['facets'][200:200+n] | |
| cosine=np.einsum('nd,nkd->nk',Q,D) | |
| base3,b_order=coverage(cosine,test,3) | |
| report=json.load(open(ROOT/('hotpot_subset_out/graph_results.json' if ds=='hotpot' else 'wiki_subset_out/graph_results.json'))) | |
| link,match=prep_links(test);docsim=np.einsum('nkd,njd->nkj',D,D); facsim=np.einsum('nmd,nkd->nmk',Fq,D) | |
| gorders=np.asarray([graph_select(cosine[i],docsim[i],facsim[i],link[i],match[i],k=5,**report['parameters']) for i in range(n)]) | |
| truth=np.zeros((n,10),np.uint8) | |
| for i,r in enumerate(test): | |
| gold={x.casefold().strip() for x,_ in r['supporting_facts']} | |
| for j,(title,_) in enumerate(r['context']):truth[i,j]=int(title.casefold().strip() in gold) | |
| cross_order=np.argsort(-scores,axis=1,kind='stable') | |
| def hit(order,k):return np.array([int(truth[i,order[i,:k]].sum()==truth[i].sum() and truth[i].sum()>0) for i in range(n)],np.int8) | |
| methods={'cosine':b_order,'cross_encoder':cross_order,'evidence_graph':gorders} | |
| metrics={name:{str(k):float(hit(order,k).mean()) for k in (1,2,3,4,5)} for name,order in methods.items()} | |
| rng=np.random.default_rng(9027);boots=rng.integers(0,n,size=(3000,n));cis={} | |
| a=hit(cross_order,3) | |
| for name,order in methods.items(): | |
| if name=='cross_encoder':continue | |
| delta=a-hit(order,3) | |
| cis[name]={'difference':float(delta.mean()),'ci95':np.quantile(delta[boots].mean(axis=1),[.025,.975]).tolist(),'wins':int((delta>0).sum()),'losses':int((delta<0).sum()),'ties':int((delta==0).sum())} | |
| result={'dataset':ds,'questions':n,'pairs':n*10,'model':'cross-encoder/ms-marco-MiniLM-L6-v2 pretrained original weights','max_wordpieces':tok.max_length,'batch_size':batch,'metrics_complete_gold_doc_coverage':metrics,'cross_encoder_minus':cis,'timing':{**perf,'total_scoring_seconds_with_tokenization':ts,'scoring_ms_per_question':ts*1000/n,'total_run_s':time.perf_counter()-start},'protocol':'Existing seeded held-out questions 200:800 (200 prior tune questions excluded). Same ten candidates. Gold labels used only to measure coverage. Graph selector retains previously tuned parameters.'} | |
| out=ROOT/f'crossencoder_{ds}_{tok.max_length}_{n}_results.json';out.write_text(json.dumps(result,indent=2)) | |
| np.savez_compressed(ROOT/f'crossencoder_{ds}_{tok.max_length}_{n}_scores.npz',scores=scores,orders=cross_order) | |
| print('RESULT',json.dumps(result,indent=2),flush=True) | |
| return result | |
| if __name__=='__main__': | |
| ap=argparse.ArgumentParser();ap.add_argument('--datasets',nargs='+',default=['hotpot','wiki']);ap.add_argument('--max-length',type=int,default=256);ap.add_argument('--batch-size',type=int,default=16);ap.add_argument('--max-queries',type=int,default=None) | |
| args=ap.parse_args();model=load_crossencoder();tokenizer=PairTokenizer(length=args.max_length) | |
| for ds in args.datasets:run(ds,model,tokenizer,args.batch_size,args.max_queries) | |