Spaces:
Sleeping
Sleeping
| """CPU-only evidence-completeness benchmark for HotpotQA distractor setting. | |
| Uses frozen local MiniLM, no gold evidence in selection. Dataset's supporting facts evaluate chosen titles. | |
| """ | |
| import os,time,re,json,sys,hashlib | |
| os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='4' | |
| from pathlib import Path | |
| import numpy as np | |
| from minilm_cpu_pipeline import load_model,WordPiece,embed | |
| from parquet_lite import load_2wiki | |
| DATA=Path('/mnt/data/validation-00000-of-00001.parquet');OUT=Path('/mnt/data/wiki_subset_out');OUT.mkdir(exist_ok=True) | |
| N_TUNE=200;N_TEST=600;RNG=20261030 | |
| def facets(q): | |
| q=q.strip().rstrip('?') | |
| # Query-provided pieces only; no gold evidence or topic annotations. | |
| p=re.split(r'(?i)\b(?:and|or|whose|which|who|that|where|while|whereas)\b|[,;:]\s*',q) | |
| p=[x.strip() for x in p if len(x.split())>=2] | |
| if len(p)<2: | |
| w=q.split();m=max(2,len(w)//2) | |
| p=[' '.join(w[:m]),' '.join(w[m:])] if len(w)>=5 else [q,q] | |
| return (p[:3]+[q]*3)[:3] | |
| def safe_key(x):return re.sub(r'\s+',' ',x.strip().lower()) | |
| def load_data(): | |
| data=[r for r in load_2wiki(str(DATA),12576) if len(r['context'])==10]; rng=np.random.default_rng(RNG) | |
| print('total 2Wiki valid candidates',len(data),flush=True) | |
| # Sample from all subtypes without using labels for feature engineering. | |
| ids=rng.permutation(len(data))[:N_TUNE+N_TEST] | |
| items=[data[int(i)] for i in ids] | |
| counts={} | |
| for r in items:counts[r['type']]=counts.get(r['type'],0)+1 | |
| print('sample',len(items),'types',counts,flush=True) | |
| return items,ids | |
| def encode_data(items): | |
| path=OUT/'wiki_minilm_800.npz' | |
| identity=hashlib.sha256('|'.join(r['_id'] for r in items).encode()).hexdigest()[:16] | |
| if path.exists(): | |
| dat=np.load(path) | |
| if str(dat['identity'])==identity: | |
| print('loaded cached embeddings',flush=True) | |
| return dat['q'],dat['docs'],dat['facets'] | |
| model=load_model();tok=WordPiece(length=128) | |
| queries=[r['question'] for r in items] | |
| docs=[title+'. '+ ' '.join(sentences) for r in items for title,sentences in r['context']] | |
| ftexts=[t for r in items for t in facets(r['question'])] | |
| print('encoding',len(queries),'queries',len(docs),'docs',len(ftexts),'aspects',flush=True) | |
| st=time.perf_counter();q=embed(queries,model,tok,batch_size=24,callback=lambda i,n,t:print('q',i,n,round(t,1),flush=True) if i>0 else None) | |
| print('query done',round(time.perf_counter()-st,1),flush=True) | |
| d=embed(docs,model,tok,batch_size=24,callback=lambda i,n,t:print('docs',i,n,round(t,1),flush=True) if i>0 else None).reshape(len(items),10,384) | |
| print('documents done',round(time.perf_counter()-st,1),flush=True) | |
| f=embed(ftexts,model,tok,batch_size=24,callback=lambda i,n,t:print('facets',i,n,round(t,1),flush=True) if i>0 else None).reshape(len(items),3,384) | |
| print('facets done',round(time.perf_counter()-st,1),flush=True) | |
| np.savez_compressed(path,identity=identity,q=q,docs=d,facets=f) | |
| return q,d,f | |
| def greedy(rel,dc,fac,method,k=5,alpha=.6,beta=.2,threshold=.35): | |
| # Single query score vector, 10x10 document similarity, facet-doc sim 3x10. | |
| K=len(rel);sel=[];remain=set(range(K)); | |
| cfac=np.clip(fac-threshold,0,None) | |
| mx=np.zeros(len(cfac));mxcluster=np.zeros(K) | |
| for t in range(k): | |
| best=(-1e9,None) | |
| for j in sorted(remain): | |
| if method=='cosine':v=rel[j] | |
| elif method=='mmr':v=rel[j] if not sel else alpha*rel[j]-(1-alpha)*max(0,max(dc[j,x] for x in sel)) | |
| elif method=='coverage':v=alpha*rel[j]+(1-alpha)*np.maximum(cfac[:,j]-mx,0).mean() | |
| elif method=='facility': | |
| weights=np.clip(rel,0,None)**2 | |
| v=alpha*rel[j]+(1-alpha)*np.dot(weights,np.maximum(np.clip(dc[:,j],0,None)-mxcluster,0))/max(weights.sum(),1e-8) | |
| elif method=='bridge': | |
| bridge=max((dc[j,x] for x in sel),default=0) | |
| v=alpha*rel[j]+(1-alpha)*bridge-beta*max(0,bridge-.8) if sel else rel[j] | |
| elif method=='combined': | |
| gain=np.maximum(cfac[:,j]-mx,0).mean() | |
| # Complementary candidate information, plus query-conditioned facet coverage. | |
| weights=np.clip(rel,0,None)**2 | |
| facility=np.dot(weights,np.maximum(np.clip(dc[:,j],0,None)-mxcluster,0))/max(weights.sum(),1e-8) | |
| v=alpha*rel[j]+(1-alpha)*(beta*gain+(1-beta)*facility) | |
| else:raise ValueError(method) | |
| if v>best[0]:best=(v,j) | |
| j=best[1];sel.append(j);remain.remove(j) | |
| mx=np.maximum(mx,cfac[:,j]);mxcluster=np.maximum(mxcluster,np.clip(dc[:,j],0,None)) | |
| return sel | |
| def eval_one(items,scores,docsim,facsim,method,params,inds,ks=(2,3,4,5)): | |
| out=np.zeros((len(inds),len(ks)),float); selected=[];goldcounts=[];omiss=[] | |
| for ii,i in enumerate(inds): | |
| r=items[i]; gold=set(safe_key(t) for t,_ in r['supporting_facts']);titles=[safe_key(x[0]) for x in r['context']] | |
| goldcounts.append(len(gold));possible=gold.issubset(set(titles));omiss.append(possible) | |
| order=greedy(scores[i],docsim[i],facsim[i],method,k=max(ks),**params) | |
| selected.append(order) | |
| for jj,k in enumerate(ks):out[ii,jj]=int(possible and gold.issubset(set(titles[x] for x in order[:k]))) | |
| return out,np.asarray(selected),np.asarray(goldcounts),np.asarray(omiss) | |
| def main(): | |
| items,orig_idx=load_data();q,d,f=encode_data(items) | |
| sims=np.einsum('nd,nkd->nk',q,d).astype('float32');dcs=np.einsum('nkd,njd->nkj',d,d).astype('float32');fs=np.einsum('nmd,nkd->nmk',f,d).astype('float32') | |
| tunes=np.arange(N_TUNE);tests=np.arange(N_TUNE,N_TUNE+N_TEST) | |
| param_grid={ | |
| 'cosine':[{}], | |
| 'mmr':[{'alpha':a} for a in [.3,.5,.7,.85,.95]], | |
| 'coverage':[{'alpha':a,'threshold':t} for a in [.1,.3,.5,.7,.9] for t in [.2,.35,.5,.65]], | |
| 'facility':[{'alpha':a} for a in [.2,.4,.6,.8,.95]], | |
| 'bridge':[{'alpha':a,'beta':b} for a in [.4,.6,.8] for b in [0,.2,.4]], | |
| 'combined':[{'alpha':a,'beta':b,'threshold':t} for a in [.2,.4,.6,.8] for b in [.25,.5,.75] for t in [.3,.5]], | |
| } | |
| best={};val={} | |
| for method,grid in param_grid.items(): | |
| st=time.perf_counter();records=[] | |
| for params in grid: | |
| arr,_,_,_=eval_one(items,sims,dcs,fs,method,params,tunes) | |
| records.append((float(arr[:,1].mean()),float(arr[:,2].mean()),params)) | |
| records.sort(key=lambda x:(x[0],x[1]),reverse=True) | |
| best[method]=records[0][2];val[method]=records[0][0] | |
| print('TUNED',method,'@3',records[0][0],best[method],'seconds',round(time.perf_counter()-st,1),flush=True) | |
| results={};test_arrays={};selections={} | |
| for method,params in best.items(): | |
| st=time.perf_counter();arr,idx,ngold,oracle=eval_one(items,sims,dcs,fs,method,params,tests) | |
| ms=(time.perf_counter()-st)/len(tests)*1000 | |
| results[method]={'parameters':params,'validation_complete_evidence_at3':val[method],'heldout_complete_evidence_at_k':{str(k):float(arr[:,j].mean()) for j,k in enumerate((2,3,4,5))},'selection_cpu_ms_per_question':ms,'gold_title_count_distribution':{str(x):int((ngold==x).sum()) for x in np.unique(ngold)},'oracle_coverage_rate':float(oracle.mean())} | |
| test_arrays[method]=arr;selections[method]=idx | |
| print('TEST',method,results[method]['heldout_complete_evidence_at_k'],'ms',round(ms,3),flush=True) | |
| base=test_arrays['cosine'];rng=np.random.default_rng(RNG) | |
| bs=rng.integers(len(tests),size=(2000,len(tests))) | |
| for method in best: | |
| delta=test_arrays[method][:,1]-base[:,1] | |
| ci=np.quantile(delta[bs].mean(axis=1),[.025,.975]).tolist() | |
| results[method]['delta_at3_vs_cosine']=float(delta.mean());results[method]['delta_at3_ci95']=ci | |
| results[method]['paired_win_loss_at3']={'wins':int((delta>0).sum()),'losses':int((delta<0).sum()),'ties':int((delta==0).sum())} | |
| ids=[items[i]['_id'] for i in tests];details=[] | |
| for j,i in enumerate(tests): | |
| r=items[i]; details.append({'id':r['_id'],'question':r['question'],'answer':r['answer'],'type':r['type'],'gold_support_titles':sorted(set(t for t,_ in r['supporting_facts'])),'cosine': [r['context'][v][0] for v in selections['cosine'][j]],'coverage':[r['context'][v][0] for v in selections['coverage'][j]],'combined':[r['context'][v][0] for v in selections['combined'][j]]}) | |
| result={'dataset':'2WikiMultiHopQA validation','source_file':str(DATA),'total_dev':12576,'sample_count':len(items),'tune_count':N_TUNE,'heldout_count':N_TEST,'random_seed':RNG,'protocol':'Each question has own ten provided paragraphs; gold supporting passage titles are evaluation only; frozen MiniLM encoded whole query, each candidate passage, and three heuristic query facets; hyperparameters chosen on 250 dev cases; tested on separate 750 dev cases. No generative LLM was invoked.','results':results} | |
| (OUT/'wiki_results.json').write_text(json.dumps(result,indent=2));(OUT/'wiki_case_studies.json').write_text(json.dumps(details,indent=2));print('DONE',flush=True) | |
| if __name__=='__main__':main() | |