Spaces:
Running on Zero
Running on Zero
| """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) | |