scripts/beir_benchmark.py

"""
BEIR SCIFACT Benchmark — naive vs full-stack RAG on a real research dataset.

Dataset:  SCIFACT (Wadden et al., 2020) — scientific claim verification
          Loaded from BeIR/scifact on HuggingFace Hub (BEIR benchmark suite)
          5,183 corpus documents, 300 test queries with gold relevance labels

Configs compared (additive stack):
  1. naive           — dense-only vector search, no reranking
  2. hybrid          — dense + BM25 + RRF fusion
  3. hybrid+rerank   — adds cross-encoder reranking
  4. full_stack      — adds HyDE query expansion

Metrics:
  Recall@K          fraction of gold-relevant docs retrieved in top-k
  Faithfulness      LLM-as-judge (Claude Sonnet): 1-5, is answer grounded in context?
  Answer Relevancy  cosine similarity between answer and question embeddings
  Latency (p50)     median end-to-end response time in ms

Usage:
    python3 scripts/beir_benchmark.py [--queries N] [--output PATH] [--skip-ingest]
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import tempfile
import time
from dataclasses import asdict, dataclass
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

logging.basicConfig(level=logging.WARNING)

from rich import box  # noqa: E402
from rich.console import Console  # noqa: E402
from rich.table import Table  # noqa: E402

console = Console()

COLLECTION_NAME = "beir_scifact"
N_DISTRACTORS = 200  # extra corpus docs to make retrieval non-trivial


# ── Data structures ───────────────────────────────────────────────────────────


@dataclass
class ConfigResult:
    name: str
    description: str
    n_queries: int
    mean_recall_at_k: float
    mean_faithfulness: float
    mean_relevancy: float
    p50_latency_ms: float
    error: str = ""


# ── Dataset loading ───────────────────────────────────────────────────────────


def load_scifact(n_queries: int) -> tuple[dict, list[dict]]:
    """
    Load SCIFACT corpus and queries from HuggingFace Hub.

    Returns:
        docs_to_ingest: {doc_id: {"title": str, "text": str}}
        qa_pairs:       list of dicts with question/relevant_sources/collection
    """
    console.print("[cyan]Loading SCIFACT from HuggingFace Hub…[/cyan]")
    from datasets import load_dataset

    corpus_ds = load_dataset("BeIR/scifact", "corpus", split="corpus")
    queries_ds = load_dataset("BeIR/scifact", "queries", split="queries")
    qrels_ds = load_dataset("BeIR/scifact-qrels", split="test")

    # Build lookup dicts
    corpus = {row["_id"]: {"title": row["title"], "text": row["text"]} for row in corpus_ds}
    queries = {row["_id"]: row["text"] for row in queries_ds}

    # qrels: query_id → set of relevant doc_ids
    qrels: dict[str, set[str]] = {}
    for row in qrels_ds:
        qid = str(row["query-id"])
        did = str(row["corpus-id"])
        if int(row["score"]) > 0:
            qrels.setdefault(qid, set()).add(did)

    # Take first N queries that have at least one relevant doc
    selected_qids = [qid for qid in queries if qid in qrels][:n_queries]
    console.print(f"  Selected {len(selected_qids)} queries with gold relevance labels")

    # Collect docs: all relevant + N_DISTRACTORS extras for realistic difficulty
    relevant_doc_ids: set[str] = set()
    for qid in selected_qids:
        relevant_doc_ids.update(qrels[qid])

    all_doc_ids = list(corpus.keys())
    distractor_ids = [d for d in all_doc_ids if d not in relevant_doc_ids][:N_DISTRACTORS]
    docs_to_ingest = {
        did: corpus[did] for did in (list(relevant_doc_ids) + distractor_ids) if did in corpus
    }

    console.print(
        f"  Corpus slice: {len(relevant_doc_ids)} relevant + {len(distractor_ids)} distractor docs = {len(docs_to_ingest)} total"
    )

    # Build QA pairs — relevant_sources uses filename format: {doc_id}.txt
    qa_pairs = [
        {
            "question": queries[qid],
            "expected_answer": "",  # SCIFACT doesn't provide reference answers
            "relevant_sources": [f"{did}.txt" for did in qrels[qid]],
            "collection": COLLECTION_NAME,
        }
        for qid in selected_qids
    ]

    return docs_to_ingest, qa_pairs


# ── Ingestion ─────────────────────────────────────────────────────────────────


def ingest_corpus(docs: dict) -> None:
    """Write corpus docs to temp files and ingest into ChromaDB."""
    from core.ingestion import ingest_document

    console.print(f"\n[cyan]Ingesting {len(docs)} documents into '{COLLECTION_NAME}'…[/cyan]")
    start = time.perf_counter()
    ok = fail = 0

    with tempfile.TemporaryDirectory() as tmpdir:
        for doc_id, doc in docs.items():
            path = Path(tmpdir) / f"{doc_id}.txt"
            content = f"{doc['title']}\n\n{doc['text']}" if doc["title"] else doc["text"]
            path.write_text(content)
            try:
                ingest_document(str(path), collection_name=COLLECTION_NAME, overwrite=False)
                ok += 1
            except Exception as e:
                fail += 1
                logging.debug("Ingest failed for %s: %s", doc_id, e)

            if (ok + fail) % 50 == 0:
                console.print(f"  {ok + fail}/{len(docs)} docs  ({ok} ok, {fail} failed)", end="\r")

    elapsed = time.perf_counter() - start
    console.print(f"  Done: {ok} ingested, {fail} failed — {elapsed:.0f}s")


# ── Per-config evaluation ─────────────────────────────────────────────────────


@dataclass
class ConfigSpec:
    name: str
    description: str
    use_hybrid: bool
    use_reranker: bool
    use_hyde: bool


CONFIGS = [
    ConfigSpec("naive", "Dense-only, no reranking", False, False, False),
    ConfigSpec("hybrid", "Dense + BM25 + RRF", True, False, False),
    ConfigSpec("hybrid+rerank", "Hybrid + cross-encoder reranking", True, True, False),
    ConfigSpec("full_stack", "Hybrid + reranking + HyDE", True, True, True),
]


def run_config(spec: ConfigSpec, qa_pairs: list[dict]) -> ConfigResult:
    """Patch settings, run all QA pairs, restore settings, return aggregate metrics."""
    from config import settings
    from core.evaluation import evaluate_sample
    from models import EvalSample

    orig_hybrid = settings.use_hybrid_search
    orig_reranker = settings.use_reranker
    orig_hyde = settings.use_hyde

    settings.use_hybrid_search = spec.use_hybrid
    settings.use_reranker = spec.use_reranker
    settings.use_hyde = spec.use_hyde

    recalls, faiths, relevancies, latencies = [], [], [], []
    error_msg = ""

    try:
        for qa in qa_pairs:
            sample = EvalSample(
                question=qa["question"],
                expected_answer=qa.get("expected_answer", ""),
                relevant_sources=qa["relevant_sources"],
                collection=qa["collection"],
            )
            try:
                r = evaluate_sample(sample)
                recalls.append(r.recall_at_k)
                faiths.append(r.faithfulness_score)
                relevancies.append(r.answer_relevancy)
                latencies.append(r.latency_ms)
            except Exception as e:
                logging.warning("Sample failed: %s", e)
    except Exception as e:
        error_msg = str(e)
    finally:
        settings.use_hybrid_search = orig_hybrid
        settings.use_reranker = orig_reranker
        settings.use_hyde = orig_hyde

    def mean(lst: list[float]) -> float:
        return round(sum(lst) / len(lst), 4) if lst else 0.0

    def median(lst: list[float]) -> float:
        if not lst:
            return 0.0
        s = sorted(lst)
        mid = len(s) // 2
        return round(s[mid] if len(s) % 2 else (s[mid - 1] + s[mid]) / 2, 1)

    return ConfigResult(
        name=spec.name,
        description=spec.description,
        n_queries=len(qa_pairs),
        mean_recall_at_k=mean(recalls),
        mean_faithfulness=mean(faiths),
        mean_relevancy=mean(relevancies),
        p50_latency_ms=median(latencies),
        error=error_msg,
    )


# ── Display ───────────────────────────────────────────────────────────────────


def print_results(results: list[ConfigResult]) -> None:
    table = Table(
        title="BEIR SCIFACT — RAG Technique Comparison",
        box=box.ROUNDED,
        header_style="bold magenta",
    )
    table.add_column("Config", style="cyan", min_width=18)
    table.add_column("Recall@K", justify="right", min_width=10)
    table.add_column("Faithfulness", justify="right", min_width=14)
    table.add_column("Relevancy", justify="right", min_width=10)
    table.add_column("Latency p50", justify="right", min_width=12)
    table.add_column("vs Naive", justify="right", min_width=22)

    baseline = results[0] if results else None

    for r in results:
        if r.error:
            table.add_row(r.name, "[red]ERROR[/red]", "—", "—", "—", r.error[:30])
            continue

        if baseline and r.name != "naive":
            rd = r.mean_recall_at_k - baseline.mean_recall_at_k
            fd = r.mean_faithfulness - baseline.mean_faithfulness
            color = "green" if (rd > 0 or fd > 0) else "yellow"
            vs = f"[{color}]recall {rd:+.3f} / faith {fd:+.2f}[/{color}]"
        else:
            vs = "[dim]baseline[/dim]"

        rc = (
            "green"
            if r.mean_recall_at_k >= 0.7
            else ("yellow" if r.mean_recall_at_k >= 0.5 else "red")
        )
        fc = (
            "green"
            if r.mean_faithfulness >= 4.0
            else ("yellow" if r.mean_faithfulness >= 3.0 else "red")
        )

        table.add_row(
            r.name,
            f"[{rc}]{r.mean_recall_at_k:.3f}[/{rc}]",
            f"[{fc}]{r.mean_faithfulness:.2f}/5.0[/{fc}]",
            f"{r.mean_relevancy:.3f}",
            f"{r.p50_latency_ms:.0f}ms",
            vs,
        )

    console.print("\n")
    console.print(table)
    n = results[0].n_queries if results else 0
    console.print(
        f"[dim]Dataset: BEIR/SCIFACT (Wadden et al., 2020) — {n} test queries, {N_DISTRACTORS + len([r for r in results if not r.error])} corpus docs[/dim]"
    )
    console.print(
        "[dim]LLM judge: Claude Sonnet (faithfulness) | Embeddings: all-MiniLM-L6-v2 (relevancy)[/dim]\n"
    )


# ── Main ──────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(description="BEIR SCIFACT RAG benchmark")
    parser.add_argument(
        "--queries", type=int, default=25, help="Number of test queries (default: 25)"
    )
    parser.add_argument(
        "--output", type=str, default="scripts/beir_results.json", help="Output JSON path"
    )
    parser.add_argument(
        "--skip-ingest", action="store_true", help="Skip ingestion (re-use existing collection)"
    )
    args = parser.parse_args()

    console.print("[bold cyan]\nBEIR SCIFACT Benchmark[/bold cyan]")
    console.print(
        f"Queries: {args.queries} | Collection: {COLLECTION_NAME} | Output: {args.output}\n"
    )

    docs, qa_pairs = load_scifact(args.queries)

    if not args.skip_ingest:
        ingest_corpus(docs)
    else:
        console.print(
            f"[yellow]Skipping ingest — using existing '{COLLECTION_NAME}' collection[/yellow]"
        )

    results: list[ConfigResult] = []
    for spec in CONFIGS:
        console.print(
            f"\n[bold]Config {len(results) + 1}/4:[/bold] {spec.name} — {spec.description}"
        )
        t0 = time.perf_counter()
        result = run_config(spec, qa_pairs)
        elapsed = time.perf_counter() - t0
        console.print(
            f"  Finished in {elapsed:.0f}s — recall={result.mean_recall_at_k:.3f}  faith={result.mean_faithfulness:.2f}  relev={result.mean_relevancy:.3f}"
        )
        results.append(result)

    print_results(results)

    output = {
        "dataset": "BEIR/SCIFACT",
        "paper": "Wadden et al. (2020) — Fact or Fiction: Verifying Scientific Claims",
        "n_queries": args.queries,
        "n_corpus_docs": len(docs),
        "llm_judge": "claude-sonnet-4-6",
        "embeddings": "all-MiniLM-L6-v2",
        "configs": [asdict(r) for r in results],
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, "w") as f:
        json.dump(output, f, indent=2)
    console.print(f"[dim]Results saved → {args.output}[/dim]")


if __name__ == "__main__":
    main()