File size: 2,799 Bytes
fd3cc1a | 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 | #!/usr/bin/env python
"""Benchmark retrieval quality across pipeline configurations.
Runs the same labelled eval set (data/eval.jsonl, 20 docs / 19 questions)
through three configurations to quantify what chunking and hybrid retrieval
actually buy you:
1. naive — whole document as one chunk, vector search only
2. + chunking — sentence-aware chunking (size=120, overlap=20), vector only
3. + hybrid — same chunking, + BM25 and Reciprocal Rank Fusion
Prints a markdown table (paste straight into the README) and writes
benchmark_results.json alongside it.
Usage:
python scripts/benchmark.py
"""
from __future__ import annotations
import json
import sys
from pathlib import Path
ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT))
from lumen_rag.embeddings import get_embedder
from lumen_rag.eval import evaluate
from lumen_rag.eval.harness import load_cases
from lumen_rag.ingestion.loaders import _LOADERS, load_file
from lumen_rag.ingestion.pipeline import ingest_documents
from lumen_rag.retrieval import Retriever
from lumen_rag.store import VectorStore
DOCS_DIR = ROOT / "data" / "docs"
EVAL_PATH = ROOT / "data" / "eval.jsonl"
K = 5
CONFIGS = [
{"name": "naive (1 chunk/doc, vector-only)", "chunk_size": 100_000, "overlap": 0, "mode": "vector"},
{"name": "+ sentence chunking (vector-only)", "chunk_size": 120, "overlap": 20, "mode": "vector"},
{"name": "+ hybrid (BM25 + RRF)", "chunk_size": 120, "overlap": 20, "mode": "hybrid"},
]
def build_retriever(chunk_size: int, overlap: int) -> Retriever:
embedder = get_embedder()
store = VectorStore(dim=embedder.dim)
docs = [load_file(p) for p in sorted(DOCS_DIR.iterdir()) if p.suffix.lower() in _LOADERS]
ingest_documents(docs, store=store, embedder=embedder, chunk_size=chunk_size, overlap=overlap)
return Retriever(store, embedder)
def main() -> None:
cases = load_cases(EVAL_PATH)
rows = []
for cfg in CONFIGS:
retriever = build_retriever(cfg["chunk_size"], cfg["overlap"])
report = evaluate(retriever, cases, k=K)
d = report.as_dict()
rows.append({"config": cfg["name"], **d})
header = ["Configuration", "recall@5", "precision@5", "MRR", "nDCG@5", "hit rate"]
lines = [
"| " + " | ".join(header) + " |",
"|" + "---|" * len(header),
]
for r in rows:
lines.append(
"| {config} | {recall@k:.2f} | {precision@k:.2f} | {mrr:.2f} | {ndcg@k:.2f} | {hit_rate:.2f} |".format(
**r
)
)
table = "\n".join(lines)
print(table)
out = ROOT / "benchmark_results.json"
out.write_text(json.dumps({"k": K, "n_cases": len(cases), "results": rows}, indent=2))
print(f"\nWrote {out}")
if __name__ == "__main__":
main()
|