src/visual_doc_rag/evaluate.py

"""Retrieval metrics using benchmark qrels; no answer quality evaluation."""

import json
import math
import platform
from datetime import UTC, datetime
from importlib.metadata import version
from pathlib import Path
from statistics import mean, median
from time import perf_counter

import torch

from visual_doc_rag.data import inspect_tables, load_tables, write_json
from visual_doc_rag.retrieve import rrf
from visual_doc_rag.text_index import TextIndex
from visual_doc_rag.visual_index import Encoder, VisualIndex


def metrics_at(ranking: list[int], qrels: dict[int, int], k: int) -> dict[str, float]:
    """Match trec_eval ndcg_cut: linear graded gains, log2(rank+1) discount."""
    if k < 1 or len(set(ranking)) != len(ranking):
        raise ValueError("Cutoff must be positive and ranking IDs must be unique.")
    relevant = {page for page, score in qrels.items() if score > 0}
    gains = [max(qrels.get(page, 0), 0) for page in ranking[:k]]
    ideal = sorted((score for score in qrels.values() if score > 0), reverse=True)[:k]
    dcg = sum(score / math.log2(rank + 2) for rank, score in enumerate(gains))
    idcg = sum(score / math.log2(rank + 2) for rank, score in enumerate(ideal))
    return {
        f"NDCG@{k}": dcg / idcg if idcg else 0.0,
        f"recall@{k}": len(set(ranking[:k]) & relevant) / len(relevant) if relevant else 0.0,
    }


def latency_summary(values: list[float]) -> dict:
    ordered = sorted(values)
    return {
        "mean_seconds": mean(values),
        "median_seconds": median(values),
        "p95_seconds": ordered[math.ceil(0.95 * len(ordered)) - 1],
    }


def evaluate(subset: str, data_dir: Path, runs_dir: Path) -> tuple[Path, dict]:
    source, tables = load_tables(subset, data_dir)
    if not all(inspect_tables(tables)["integrity_checks"].values()):
        raise ValueError("Dataset integrity failed. Run `vdr inspect` before evaluating.")
    root = data_dir / subset / "indexes"
    text = TextIndex.load(root, source)
    visual = VisualIndex(root, source)
    expected_ids = set(tables["corpus"]["corpus_id"])
    if set(text.page_ids) != expected_ids or set(visual.page_ids) != expected_ids:
        raise ValueError("Indexes must cover the full candidate corpus.")
    encoder = Encoder(data_dir, visual.manifest["signature"]["model"], "cpu")
    queries = [
        row for row in tables["queries"] if row["language"].strip().lower() in {"english", "en"}
    ]
    qrels: dict[int, dict[int, int]] = {}
    for row in tables["qrels"]:
        qrels.setdefault(row["query_id"], {})[row["corpus_id"]] = row["score"]
    # Warm model kernels and search; exclude initialization/warmup from measured latency.
    visual.search(encoder.encode(query=queries[0]["query"]))
    text.search(queries[0]["query"])
    timings = {name: [] for name in ("text", "visual", "hybrid")}
    per_query = []
    stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%S.%fZ")
    output = runs_dir / stamp
    output.mkdir(parents=True)
    for index, row in enumerate(queries):
        started = perf_counter()
        text_ranking = text.search(row["query"])
        text_seconds = perf_counter() - started
        started = perf_counter()
        visual_ranking = visual.search(encoder.encode(query=row["query"]))
        visual_seconds = perf_counter() - started
        started = perf_counter()
        hybrid_ranking = rrf([text_ranking, visual_ranking], k=60)
        fusion_seconds = perf_counter() - started
        timings["text"].append(text_seconds)
        timings["visual"].append(visual_seconds)
        timings["hybrid"].append(text_seconds + visual_seconds + fusion_seconds)
        results = {}
        for name, ranking in zip(
            timings, (text_ranking, visual_ranking, hybrid_ranking), strict=True
        ):
            results[name] = {
                **metrics_at(ranking, qrels[row["query_id"]], 5),
                **metrics_at(ranking, qrels[row["query_id"]], 10),
            }
        per_query.append({"query_id": row["query_id"], "metrics": results})
        if (index + 1) % 10 == 0 or index + 1 == len(queries):
            print(f"Evaluated: {index + 1}/{len(queries)} English queries", flush=True)
    keys = ("NDCG@5", "NDCG@10", "recall@5", "recall@10")
    report = {
        "dataset": source,
        "n_queries": len(queries),
        "n_pages": len(expected_ids),
        "language": "english",
        "split": "test",
        "tuning": "none",
        "metrics": {
            name: {key: mean(row["metrics"][name][key] for row in per_query) for key in keys}
            for name in timings
        },
        "cpu_latency": {name: latency_summary(values) for name, values in timings.items()},
        "latency_protocol": "4 torch threads; CPU float32 query encoding + full-corpus search; "
        "warm indexes/model; no query cache; hybrid sequential sum plus RRF",
        "metric_definition": "macro mean over English queries; trec_eval linear graded NDCG; "
        "recall uses score>0; unjudged gain=0; RRF k=60 over full rankings",
        "index_build": {
            "text": json.loads((root / "text-build.json").read_text(encoding="utf-8")),
            "visual": {k: v for k, v in visual.manifest.items() if k not in {"pages", "dataset"}},
        },
        "environment": {
            "python": platform.python_version(),
            "platform": platform.platform(),
            "torch": torch.__version__,
            "torch_threads": torch.get_num_threads(),
            "processor": platform.processor(),
            "library_versions": {
                name: version(name)
                for name in ("colpali-engine", "transformers", "rank-bm25", "datasets")
            },
        },
    }
    write_json(output / "per-query.json", {"queries": per_query})
    write_json(output / "metrics.json", report)
    return output, report


def format_table(report: dict) -> str:
    header = "method     NDCG@5   NDCG@10  recall@5  recall@10"
    lines = [header]
    for name, values in report["metrics"].items():
        lines.append(f"{name:<10} " + "  ".join(f"{value:.4f}" for value in values.values()))
    return "\n".join(lines)