194 lines
7.5 KiB
Python
194 lines
7.5 KiB
Python
"""Run all retrievers against eval/queries.jsonl, emit a markdown report.
|
|
|
|
Metrics computed per retriever:
|
|
|
|
MRR — mean reciprocal rank of the FIRST expected page in the
|
|
ranked result list (0 if not in top-k).
|
|
Recall@K — fraction of expected pages that appear in top-K.
|
|
nDCG@K — discounted gain weighted by rank position.
|
|
|
|
The "right" number depends on what you're measuring. MRR tracks "the
|
|
first-line answer is correct"; Recall@K tracks "everything relevant
|
|
is there to draw from"; nDCG@K is a smoother combination of both.
|
|
For docs-RAG, MRR is usually the headline metric.
|
|
|
|
Usage:
|
|
|
|
python -m eval.run_eval \\
|
|
--queries eval/queries.jsonl \\
|
|
--k 5 \\
|
|
--output eval/results/baseline.md
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
|
|
def load_queries(path: Path) -> list[dict]:
|
|
with open(path) as fh:
|
|
return [json.loads(line) for line in fh if line.strip()]
|
|
|
|
|
|
def p_at_1(retrieved: list[tuple[str, str]], expected: list[tuple[str, str]]) -> float:
|
|
if not retrieved or not expected:
|
|
return 0.0
|
|
return 1.0 if retrieved[0] in set(expected) else 0.0
|
|
|
|
|
|
def reciprocal_rank(retrieved: list[tuple[str, str]], expected: list[tuple[str, str]]) -> float:
|
|
expected_set = set(expected)
|
|
for i, page in enumerate(retrieved, start=1):
|
|
if page in expected_set:
|
|
return 1.0 / i
|
|
return 0.0
|
|
|
|
|
|
def recall_at_k(retrieved: list[tuple[str, str]], expected: list[tuple[str, str]], k: int) -> float:
|
|
if not expected:
|
|
return 0.0
|
|
retrieved_set = set(retrieved[:k])
|
|
hits = sum(1 for e in expected if e in retrieved_set)
|
|
return hits / len(expected)
|
|
|
|
|
|
def ndcg_at_k(retrieved: list[tuple[str, str]], expected: list[tuple[str, str]], k: int) -> float:
|
|
expected_set = set(expected)
|
|
dcg = 0.0
|
|
for i, page in enumerate(retrieved[:k], start=1):
|
|
if page in expected_set:
|
|
dcg += 1.0 / math.log2(i + 1)
|
|
# Ideal DCG: every expected page in the top positions.
|
|
idcg = sum(1.0 / math.log2(i + 1) for i in range(1, min(len(expected), k) + 1))
|
|
return dcg / idcg if idcg else 0.0
|
|
|
|
|
|
def main() -> int:
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--queries", type=Path, default=Path("eval/queries.jsonl"))
|
|
p.add_argument("--k", type=int, default=5)
|
|
p.add_argument("--ks", default="1,5,10,20", help="comma-separated k-curve")
|
|
p.add_argument("--output", type=Path, default=Path("eval/results/baseline.md"))
|
|
p.add_argument("--compare", type=Path, default=None)
|
|
args = p.parse_args()
|
|
ks = sorted({int(x) for x in args.ks.split(",") if x.strip()}) or [args.k]
|
|
max_k = max(ks + [args.k])
|
|
|
|
if not args.queries.exists():
|
|
print(f"queries file not found: {args.queries}")
|
|
print("hint: copy eval/queries.jsonl.example and edit")
|
|
return 1
|
|
|
|
queries = load_queries(args.queries)
|
|
print(f"loaded {len(queries)} queries")
|
|
|
|
import os
|
|
import chromadb
|
|
from chromadb.config import Settings
|
|
from rag.embeddings import embedding_function
|
|
from rag.bm25 import BM25Index
|
|
from eval.retrievers import DenseRetriever, BM25Retriever, HybridRetriever
|
|
|
|
product = os.environ.get("PRODUCT_NAME", "morpheus")
|
|
repo_root = Path(__file__).resolve().parent.parent
|
|
client = chromadb.PersistentClient(path=str(repo_root / "chroma"),
|
|
settings=Settings(anonymized_telemetry=False))
|
|
col = client.get_collection(f"{product}_docs", embedding_function=embedding_function())
|
|
bm = BM25Index(str(repo_root / "bm25" / f"{product}_docs.db"))
|
|
|
|
from eval.retrievers import RerankedRetriever
|
|
|
|
dense = DenseRetriever(col)
|
|
bm25 = BM25Retriever(bm)
|
|
hybrid = HybridRetriever(DenseRetriever(col, pool=100), BM25Retriever(bm, pool=100))
|
|
|
|
retrievers = [dense, bm25, hybrid]
|
|
|
|
rerank_url = os.environ.get("RERANK_URL", "").rstrip("/")
|
|
if rerank_url:
|
|
retrievers += [
|
|
RerankedRetriever(bm25, col, rerank_url, name_suffix="rerank", pool=50),
|
|
RerankedRetriever(hybrid, col, rerank_url, name_suffix="rerank", pool=50),
|
|
]
|
|
print(f"reranker enabled: {rerank_url}")
|
|
|
|
rows: dict[str, dict[str, float]] = {}
|
|
per_query: list[dict] = []
|
|
for r in retrievers:
|
|
mrr_sum = recall_sum = ndcg_sum = p1_sum = 0.0
|
|
elapsed_sum = 0.0
|
|
for q in queries:
|
|
expected = [(e["bundle_id"], e["page_id"]) for e in q["expected"]]
|
|
t0 = time.time()
|
|
retrieved = r.retrieve(q["query"], k=max(max_k, 10))
|
|
elapsed = time.time() - t0
|
|
mrr = reciprocal_rank(retrieved, expected)
|
|
p1 = p_at_1(retrieved, expected)
|
|
recall = recall_at_k(retrieved, expected, args.k)
|
|
ndcg = ndcg_at_k(retrieved, expected, args.k)
|
|
mrr_sum += mrr
|
|
p1_sum += p1
|
|
recall_sum += recall
|
|
ndcg_sum += ndcg
|
|
elapsed_sum += elapsed
|
|
per_query.append({
|
|
"retriever": r.name, "query": q["query"],
|
|
"mrr": mrr, "p_at_1": int(p1), "recall@k": recall, "ndcg@k": ndcg,
|
|
"top1": list(retrieved[0]) if retrieved else None,
|
|
"ranked": [list(p) for p in retrieved],
|
|
"elapsed_s": round(elapsed, 3),
|
|
})
|
|
n = len(queries)
|
|
rows[r.name] = {
|
|
"P@1": p1_sum / n,
|
|
"MRR": mrr_sum / n,
|
|
f"Recall@{args.k}": recall_sum / n,
|
|
f"nDCG@{args.k}": ndcg_sum / n,
|
|
"avg_latency_s": elapsed_sum / n,
|
|
}
|
|
print(f" {r.name}: P@1={rows[r.name]['P@1']:.3f} "
|
|
f"MRR={rows[r.name]['MRR']:.3f} "
|
|
f"Recall@{args.k}={rows[r.name][f'Recall@{args.k}']:.3f} "
|
|
f"nDCG@{args.k}={rows[r.name][f'nDCG@{args.k}']:.3f} "
|
|
f"avg={rows[r.name]['avg_latency_s']*1000:.0f}ms")
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
md = [f"# Retrieval eval — k={args.k}", "",
|
|
f"_{len(queries)} hand-curated queries, generated {time.strftime('%Y-%m-%d %H:%M:%S')}_", "",
|
|
"| Retriever | P@1 | MRR | Recall@{k} | nDCG@{k} | avg latency |".replace("{k}", str(args.k)),
|
|
"| --- | ---: | ---: | ---: | ---: | ---: |"]
|
|
for name, m in rows.items():
|
|
md.append(f"| `{name}` | {m['P@1']:.3f} | {m['MRR']:.3f} | {m[f'Recall@{args.k}']:.3f} "
|
|
f"| {m[f'nDCG@{args.k}']:.3f} | {m['avg_latency_s']*1000:.0f}ms |")
|
|
md += ["", "## Per-query results", "",
|
|
"| Retriever | Query | MRR | top-1 |", "| --- | --- | ---: | --- |"]
|
|
for r in per_query:
|
|
top1 = f"`{r['top1'][0]}/{r['top1'][1][:24]}...`" if r["top1"] else "—"
|
|
md.append(f"| `{r['retriever']}` | {r['query'][:60]} | {r['mrr']:.3f} | {top1} |")
|
|
args.output.write_text("\n".join(md) + "\n")
|
|
sidecar = args.output.with_suffix(".jsonl")
|
|
with open(sidecar, "w") as fh:
|
|
for r in per_query:
|
|
fh.write(json.dumps({
|
|
"query": r["query"],
|
|
"retriever": r["retriever"],
|
|
"ranked": r.get("ranked") or [],
|
|
"rr": r["mrr"],
|
|
"p_at_1": r["p_at_1"],
|
|
}) + "\n")
|
|
print(f"wrote {args.output}")
|
|
print(f"wrote {sidecar}")
|
|
if args.compare:
|
|
from eval.pvalue import compare, load_sidecar, render
|
|
print()
|
|
print(render(compare(load_sidecar(sidecar), load_sidecar(args.compare))), end="")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|