"""Evaluate real retrieval on the sampled FEVER and SciFact claims.""" from __future__ import annotations import argparse from dataclasses import dataclass from pathlib import Path from statistics import mean from time import perf_counter from evaluation.reporting import write_report from evaluation.sample_benchmarks import load_evidence_corpus, load_records, relevant_doc_ids from retrieval.bm25 import BM25Retriever from retrieval.dense import DenseRetriever, load_embedder from retrieval.hybrid import HybridRetriever from retrieval.metrics import mean_reciprocal_rank, ndcg_at_k, recall_at_k @dataclass(frozen=True) class RetrievalMetrics: recall_at_1: float recall_at_5: float recall_at_10: float mrr: float ndcg_at_10: float def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Run retrieval evaluation on sampled datasets.") parser.add_argument("--data-dir", default="data/processed") parser.add_argument("--split", choices=["train", "val", "test"], default="val") parser.add_argument("--max-queries", type=int, default=20) parser.add_argument("--dense-backend", choices=["hashing", "sentence-transformers"], default="hashing") parser.add_argument("--embedding-model", default="sentence-transformers/all-MiniLM-L6-v2") parser.add_argument("--top-k", default="1,5,10", help="Comma-separated retrieval cutoffs, e.g. 1,5,10") parser.add_argument("--output-json", default="reports/retrieval_eval.json") parser.add_argument("--output-md", default="reports/retrieval_eval.md") return parser def main() -> None: # pragma: no cover - script entrypoint args = build_parser().parse_args() data_dir = Path(args.data_dir) output_json = Path(args.output_json) output_md = Path(args.output_md) output_json.parent.mkdir(parents=True, exist_ok=True) output_md.parent.mkdir(parents=True, exist_ok=True) try: report = evaluate_retrieval( data_dir=data_dir, split=args.split, max_queries=args.max_queries, dense_backend=args.dense_backend, embedding_model=args.embedding_model, top_k=parse_top_k(args.top_k), ) except Exception as exc: # pragma: no cover - failure path is surfaced explicitly failure_path = output_md.with_name(f"{output_md.stem}_FAILED.md") failure_path.write_text( "\n".join( [ "# Retrieval Evaluation Failed", "", f"- split: {args.split}", f"- dense_backend: {args.dense_backend}", f"- embedding_model: {args.embedding_model}", "", "The retrieval evaluation could not complete.", "", f"Reason: {type(exc).__name__}: {exc}", "", "No metrics were fabricated.", ] ), encoding="utf-8", ) raise SystemExit(f"Retrieval evaluation failed: {exc}") from exc write_report(report, output_json) output_md.write_text(_to_markdown(report), encoding="utf-8") print(f"Wrote retrieval evaluation to {output_json}") def evaluate_retrieval( *, data_dir: Path, split: str, max_queries: int, dense_backend: str, embedding_model: str, top_k: list[int], ) -> dict[str, object]: started = perf_counter() records_by_split = load_records(data_dir) corpus = load_evidence_corpus(data_dir) records = _select_records(records_by_split, split)[:max_queries] if not records: raise ValueError(f"No records found for split={split!r}") embedder = load_embedder( "sentence-transformers" if dense_backend == "sentence-transformers" else "hashing", embedding_model if dense_backend == "sentence-transformers" else None, allow_fallback=False, ) bm25 = BM25Retriever(corpus) dense = DenseRetriever(corpus, embedder=embedder) hybrid = HybridRetriever(bm25, dense) max_cutoff = max(top_k) rankings = { "bm25": [bm25.retrieve(record.claim, top_k=max_cutoff) for record in records], "dense": [dense.retrieve(record.claim, top_k=max_cutoff) for record in records], "hybrid": [hybrid.retrieve(record.claim, top_k=max_cutoff) for record in records], } relevant_sets = [relevant_doc_ids(record) for record in records] metrics = {name: _summarize(rankings[name], relevant_sets, top_k) for name in rankings} report = { "split": split, "max_queries": max_queries, "num_queries": len(records), "evidence_corpus_size": len(corpus), "dense_backend": dense.backend_name, "embedding_model": getattr(embedder, "model_name", embedding_model if dense_backend == "sentence-transformers" else "hashing"), "top_k": top_k, "runtime_seconds": round(perf_counter() - started, 3), "metrics": metrics, "limitations": _limitations(dense_backend, dense.backend_name, len(records), max_queries), } return report def _summarize( rankings: list[list[object]], relevant_sets: list[list[str]], top_k: list[int], ) -> dict[str, float]: doc_ids = [[getattr(span, "doc_id", "") for span in ranking] for ranking in rankings] return { "recall@1": mean(recall_at_k(ids, relevant, 1) for ids, relevant in zip(doc_ids, relevant_sets)), "recall@5": mean(recall_at_k(ids, relevant, 5) for ids, relevant in zip(doc_ids, relevant_sets)), "recall@10": mean(recall_at_k(ids, relevant, 10) for ids, relevant in zip(doc_ids, relevant_sets)), "mrr": mean_reciprocal_rank(doc_ids, relevant_sets), "ndcg@10": mean(ndcg_at_k(ids, relevant, 10) for ids, relevant in zip(doc_ids, relevant_sets)), } def _limitations(requested_backend: str, actual_backend: str, num_queries: int, max_queries: int) -> list[str]: limitations = [ f"Evaluated on a sample of {num_queries} queries, capped at --max-queries={max_queries}.", ] if requested_backend == "hashing": limitations.append("Dense retrieval used the hashing backend, which is a lightweight baseline.") if requested_backend == "sentence-transformers" and actual_backend != "sentence-transformers": limitations.append("Requested neural backend was not available and fell back to hashing.") return limitations def _select_records(records_by_split: dict[str, list[object]], split: str): selected = [] for split_name, split_records in records_by_split.items(): if split_name.endswith(f"_{split}"): selected.extend(split_records) return selected def parse_top_k(value: str) -> list[int]: cuts = [] for item in value.split(","): item = item.strip() if not item: continue cuts.append(int(item)) if not cuts: raise ValueError("At least one top-k cutoff is required") return sorted(set(cuts)) def _to_markdown(report: dict[str, object]) -> str: metrics = report["metrics"] lines = [ "# Retrieval Evaluation", "", f"- Split: {report['split']}", f"- Max queries: {report['max_queries']}", f"- Queries evaluated: {report['num_queries']}", f"- Evidence corpus size: {report['evidence_corpus_size']}", f"- Dense backend: {report['dense_backend']}", f"- Embedding model: {report['embedding_model']}", f"- Runtime seconds: {report['runtime_seconds']}", "", "| retriever | recall@1 | recall@5 | recall@10 | mrr | ndcg@10 |", "| --- | --- | --- | --- | --- | --- |", ] for retriever_name in ("bm25", "dense", "hybrid"): score = metrics[retriever_name] lines.append( "| " + " | ".join( [ retriever_name, f"{score['recall@1']:.3f}", f"{score['recall@5']:.3f}", f"{score['recall@10']:.3f}", f"{score['mrr']:.3f}", f"{score['ndcg@10']:.3f}", ] ) + " |" ) lines.extend(["", "## Limitations", ""]) for item in report["limitations"]: lines.append(f"- {item}") return "\n".join(lines) + "\n" if __name__ == "__main__": # pragma: no cover - script entrypoint main()