hvm-docs/eval/run_eval.py

"""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 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("--output", type=Path, default=Path("eval/results/baseline.md"))
    args = p.parse_args()

    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", "hvm")
    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 = 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(args.k, 10))
            elapsed = time.time() - t0
            mrr = reciprocal_rank(retrieved, expected)
            recall = recall_at_k(retrieved, expected, args.k)
            ndcg = ndcg_at_k(retrieved, expected, args.k)
            mrr_sum += mrr
            recall_sum += recall
            ndcg_sum += ndcg
            elapsed_sum += elapsed
            per_query.append({
                "retriever": r.name, "query": q["query"],
                "mrr": mrr, "recall@k": recall, "ndcg@k": ndcg,
                "top1": list(retrieved[0]) if retrieved else None,
                "elapsed_s": round(elapsed, 3),
            })
        n = len(queries)
        rows[r.name] = {
            "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}: 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 | MRR | Recall@{k} | nDCG@{k} | avg latency |".replace("{k}", str(args.k)),
          "| --- | ---: | ---: | ---: | ---: |"]
    for name, m in rows.items():
        md.append(f"| `{name}` | {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")
    print(f"wrote {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())