Datasetpublic

Angshul/SparseGeometricRAG

SparseGeometricRAG CPU-first sparse geometric retrieval for practical top-10 RAG No transformer inference at retrieval time. No retrieval GPU requirement. No dense document-vector dot products. No external API. SparseGeometricRAG is a retrieval system built around one systems objective: make the retrieval layer cheap enough to run on ordinary multicore CPU hardware without turning the corpus into a dense embedding database. It uses sparse TF-IDF geometry, fuzzy… See the full description on the dataset page: https://huggingface.co/datasets/Angshul/SparseGeometricRAG.

sourceHugging Facefair-noncommercial-research-licenseupdated 2mo agoView on Hugging Face

0likes331downloads

rootexperimentsmsmarco_scalemsmarco_preselection_finalists_validation.py

msmarco_preselection_finalists_validation.py59 linesDownload Raw Back to msmarco_scale

1from __future__ import annotations2import sys,time,json3import numpy as np,pandas as pd4from numba import njit,prange,set_num_threads5sys.path.insert(0,'/mnt/data')6import msmarco_best_tail_core as b7import msmarco_full_search_uniform1m as m8ROOT=m.ROOT; WORK=m.WORK; idx=b.idx; M=m.M; P=2000; set_num_threads(5)9CONFIGS=[('base',.2,1.),('b15e4',.15,4.),('b20e3',.2,3.),('b20e4',.2,4.)]10FINAL_B=.1; ALPHA=.25; WLEX=4.; WSEM=.311 12def topk(score,k):13 n=len(score); k=min(k,n); ii=np.argpartition(score,-k)[-k:] if n>k else np.arange(n); return ii[np.argsort(score[ii])[::-1]]14@njit(parallel=True,cache=False)15def finalfeat(dd,ip,ids,lexvec,semvec,dl,avgdl):16 n=len(dd); lx=np.zeros(n,np.float32); sm=np.zeros(n,np.float32); qc=np.zeros(n,np.float32)17 for z in prange(n):18  d=int(dd[z]); a=int(ip[d]); bb=int(ip[d+1]); raw=0.; s=0.; c=0.19  for kk in range(a,bb):20   t=int(ids[kk]); v=lexvec[t]21   if v>0: raw+=v; c+=122   s+=semvec[t]23  r=float(dl[d])/avgdl; den=(1-FINAL_B)+FINAL_B*r; lx[z]=raw/den; sm[z]=s; qc[z]=c24 return lx,sm,qc25 26def prep(text):27 q=idx.query_vec(text); qd=np.zeros(M,np.float32); qd[q.indices]=q.data; rterms,rd=idx.route(q); spans=[(int(j),int(idx.offs[j]),int(idx.offs[j+1])) for j in rterms if idx.offs[j+1]>idx.offs[j]]28 if not spans:return None29 docs=np.concatenate([np.asarray(idx.pd[a:bb]) for j,a,bb in spans]).astype(np.uint32,copy=False); mm=np.concatenate([np.asarray(idx.pm[a:bb]) for j,a,bb in spans]).astype(np.float32,copy=False); rt=np.concatenate([np.asarray(idx.pr[a:bb]) for j,a,bb in spans]).astype(np.uint16,copy=False); sb=np.concatenate([np.asarray(idx.ps[a:bb]) for j,a,bb in spans]).astype(np.uint16,copy=False)30 nr=len(spans); cent=np.zeros((nr,M),np.float32); rel=np.zeros((nr,M),np.float32); rho=np.empty(nr,np.float32)31 for u,(j,a,bb) in enumerate(spans):32  row=np.asarray(idx.ct[j]); ok=row!=65535; tids=row[ok].astype(np.int32,copy=False); cent[u,tids]=np.asarray(idx.cv[j])[ok]; ra=int(idx.rp[j]); rb=int(idx.rp[j+1]); rel[u,np.asarray(idx.ri[ra:rb],np.int32)]=np.asarray(idx.rv[ra:rb]); rho[u]=rd[j]33 rslot=np.concatenate([np.full(bb-a,u,dtype=np.uint8) for u,(j,a,bb) in enumerate(spans)]); base,sig,cons=b.score_components(rslot,mm,rt,sb,qd,rho,cent,rel); ud,inv=np.unique(docs,return_inverse=True); tail=np.bincount(inv,weights=base*np.power(sig,b.GAMMA,dtype=np.float32),minlength=len(ud)).astype(np.float32)+b.LAM*np.bincount(inv,weights=cons,minlength=len(ud)).astype(np.float32)34 lexvec=np.zeros(M,np.float32); lexvec[q.indices]=idx.idf[q.indices]; zero=np.zeros(M,np.float32); lex02,_=m.score_support_pool(ud,idx.sup_ip,idx.sup_ids,lexvec,zero,idx.dl,idx.avgdl); ratio=np.asarray(idx.dl[ud],np.float32)/idx.avgdl; raw=lex02*((1-.2)+.2*ratio)35 semvec=np.zeros(M,np.float32)36 for t,amp in zip(q.indices,q.data):37  a,bb=idx.A.indptr[t],idx.A.indptr[t+1]; nb=idx.A.indices[a:bb][:m.SEMK]; sv=idx.A.data[a:bb][:m.SEMK]; semvec[nb]+=float(amp)*sv*idx.idf[nb]38 return q,ud,tail,raw,ratio,lexvec,semvec39 40def rank_cfg(p,preb,eta):41 q,ud,tail,raw,ratio,lexvec,semvec=p; prelex=raw/np.maximum((1-preb)+preb*ratio,1e-6); sel=topk(m.zscore(tail)+eta*m.zscore(prelex),P); dd=ud[sel]; lx,sm,qc=finalfeat(dd,idx.sup_ip,idx.sup_ids,lexvec,semvec,idx.dl,idx.avgdl); cov=qc/max(1,len(q.indices)); ladj=lx*np.power(np.maximum(cov,1e-6),ALPHA); fin=m.zscore(tail[sel])+WLEX*m.zscore(ladj)+WSEM*m.zscore(sm); oo=np.argsort(fin)[::-1][:100]; return [int(x) for x in dd[oo]],sel42 43tr=pd.read_csv(ROOT/'train.tsv',sep='\t',usecols=['query-id']); uq=np.unique(tr['query-id'].to_numpy()); del tr; rng=np.random.default_rng(20260815); qids=[str(x) for x in rng.choice(uq,size=1000,replace=False)]; texts=m.load_query_texts(qids); qrels=m.qrels_from_tsv(ROOT/'train.tsv',qids,positive_only=True)44_=finalfeat(np.array([0],np.uint32),idx.sup_ip,idx.sup_ids,np.zeros(M,np.float32),np.zeros(M,np.float32),idx.dl,idx.avgdl); _=prep(texts[qids[0]])45runs={n:{} for n,_,_ in CONFIGS}; hits={n:0 for n,_,_ in CONFIGS}; den=0; times=[]46for zi,qid in enumerate(qids):47 t=time.perf_counter(); p=prep(texts[qid]); times.append((time.perf_counter()-t)*1000); rels=[int(d) for d,r in qrels[qid].items() if r>0]; den+=len(rels)48 if p is None:49  for n,_,_ in CONFIGS:runs[n][qid]=[]50  continue51 ud=p[1]52 for n,bb,eta in CONFIGS:53  rank,sel=rank_cfg(p,bb,eta); runs[n][qid]=rank; pool=set(map(int,ud[sel].tolist())); hits[n]+=sum(d in pool for d in rels)54 if (zi+1)%100==0:print('q',zi+1,'medianprep',float(np.median(times)),flush=True)55rows=[]56for n,bb,eta in CONFIGS:57 met=m.eval_run(runs[n],qrels); row={'name':n,'pre_b':bb,'eta':eta,'pool_relevant_recall':hits[n]/den,**met}; rows.append(row); print(row,flush=True)58rows.sort(key=lambda r:(r['nDCG@10'],r['MRR@10'],r['R@100']),reverse=True); out={'protocol':'preselection finalists selected from pool-survival validation, then structural final rule b=.1 alpha=.25 wl4 ws.3 evaluated on same deterministic TRAIN validation','rows':rows,'best':rows[0]}; json.dump(out,open(WORK/'preselection_finalists_validation.json','w'),indent=2); print('BEST',rows[0])59