Spaces:
Sleeping
Sleeping
File size: 7,460 Bytes
dd990bf | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """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
@torch.inference_mode()
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)
|