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.
0331
1from __future__ import annotations2import argparse3import json4from pathlib import Path5 6from .beir import load_beir_zip, load_beir_directory7from .config import FrozenConfig8from .index import GeometricIndex9from .metrics import evaluate_run10 11 12def _dataset(path: str, split: str):13 return load_beir_zip(path, split) if str(path).lower().endswith(".zip") else load_beir_directory(path, split)14 15 16def cmd_build(args):17 ds = _dataset(args.dataset, args.split)18 cfg = FrozenConfig(max_features=args.max_features, min_df=args.min_df)19 idx = GeometricIndex.build(ds.corpus_texts, ds.corpus_ids, cfg, verbose=True)20 idx.save(args.output)21 print(f"saved index -> {args.output}")22 23 24def cmd_eval(args):25 ds = _dataset(args.dataset, args.split)26 idx = GeometricIndex.load(args.index)27 # Evaluate only qrels-bearing queries.28 queries = {qid: ds.queries[qid] for qid in ds.qrels if qid in ds.queries}29 run, timing = idx.batch_search(queries, k=args.k, timing=True)30 metrics = evaluate_run(run, ds.qrels, ks=(10, 100), ndcg_k=10, mrr_k=10)31 out = {"dataset": ds.name, **metrics, **timing}32 print(json.dumps(out, indent=2, sort_keys=True))33 if args.run_json:34 Path(args.run_json).write_text(json.dumps(run, indent=1))35 36 37def cmd_search(args):38 idx = GeometricIndex.load(args.index)39 ids, scores = idx.search(args.query, k=args.k, return_scores=True)40 for r, (d, s) in enumerate(zip(ids, scores), start=1):41 print(f"{r:3d}\t{d}\t{s:.6f}")42 43 44def main():45 p = argparse.ArgumentParser(prog="geomretrieval")46 sp = p.add_subparsers(dest="cmd", required=True)47 48 b = sp.add_parser("build", help="Build frozen sparse index from a BEIR dataset/archive")49 b.add_argument("dataset")50 b.add_argument("output")51 b.add_argument("--split", default="test")52 b.add_argument("--max-features", type=int, default=50_000)53 b.add_argument("--min-df", type=int, default=1)54 b.set_defaults(func=cmd_build)55 56 e = sp.add_parser("eval", help="Evaluate an existing index on BEIR qrels")57 e.add_argument("dataset")58 e.add_argument("index")59 e.add_argument("--split", default="test")60 e.add_argument("--k", type=int, default=100)61 e.add_argument("--run-json", default=None)62 e.set_defaults(func=cmd_eval)63 64 s = sp.add_parser("search", help="Search an existing index")65 s.add_argument("index")66 s.add_argument("query")67 s.add_argument("--k", type=int, default=10)68 s.set_defaults(func=cmd_search)69 70 args = p.parse_args()71 args.func(args)72 73 74if __name__ == "__main__":75 main()76