Add T²-RAGBench FinQA evaluation dataset

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `t2_finqa` evaluation dataset (T²-RAGBench FinQA subset, `G4KMU/t2-ragbench`): 2,789 single-page PDFs / 8,281 numeric QA, ingested via docling with `uri = context_id` and gold retrieval keyed on `context_id`.
+
 ### Fixed
 
 - Skill tools (`search`/`cite`/`list_documents`/`get_document`) and the analysis sandbox serialize access to the shared LanceDB connection through one lock, so a turn's concurrently executed tool calls no longer trigger `RuntimeError: Already borrowed`.

evaluations/evaluations/datasets/__init__.py

@@ -1,6 +1,7 @@
 from evaluations.config import DatasetSpec
 
 from .open_rag_bench import ORB_MULTIMODAL_SPEC, ORB_TEXT_SPEC
+from .t2_ragbench import T2_FINQA_SPEC
 from .wix import WIX_SPEC
 
 DATASETS: dict[str, DatasetSpec] = {
@@ -9,6 +10,7 @@ DATASETS: dict[str, DatasetSpec] = {
         WIX_SPEC,
         ORB_TEXT_SPEC,
         ORB_MULTIMODAL_SPEC,
+        T2_FINQA_SPEC,
     )
 }
 

evaluations/evaluations/datasets/t2_ragbench.py

@@ -0,0 +1,144 @@
+import json
+import shutil
+from collections.abc import Mapping
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset
+from huggingface_hub import hf_hub_download
+from pydantic_evals import Case
+
+from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
+from evaluations.evaluators import MAPEvaluator
+
+REPO_ID = "G4KMU/t2-ragbench"
+SPLITS = ("dev", "test", "train")
+
+
+def get_cache_dir() -> Path:
+    cache_dir = Path.home() / ".cache" / "haiku.rag" / "evaluations" / "t2_pdfs"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+_rows_cache: dict[str, list[dict[str, Any]]] = {}
+
+
+def _load_rows(subset: str) -> list[dict[str, Any]]:
+    cached = _rows_cache.get(subset)
+    if cached is not None:
+        return cached
+
+    rows: list[dict[str, Any]] = []
+    for split in SPLITS:
+        path = hf_hub_download(
+            REPO_ID,
+            f"data/{subset}/{split}/metadata.jsonl",
+            repo_type="dataset",
+        )
+        with open(path) as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                rows.append({**json.loads(line), "subset": subset})
+
+    _rows_cache[subset] = rows
+    return rows
+
+
+def download_t2_pdf(subset: str, split: str, file_name: str) -> Path:
+    # hf_hub_download may return a content-addressed blob path with no suffix;
+    # the converter dispatches on extension, so materialize a real ``.pdf`` file.
+    dest = get_cache_dir() / f"{subset}_{split}_{file_name.replace('/', '_')}"
+    if dest.exists():
+        return dest
+
+    src = hf_hub_download(
+        REPO_ID,
+        f"data/{subset}/{split}/{file_name}",
+        repo_type="dataset",
+    )
+    shutil.copyfile(src, dest)
+    return dest
+
+
+def load_t2_corpus(subset: str) -> Dataset:
+    seen: set[str] = set()
+    docs: list[dict[str, Any]] = []
+    for row in _load_rows(subset):
+        context_id = row["context_id"]
+        if context_id in seen:
+            continue
+        seen.add(context_id)
+        docs.append(row)
+    return Dataset.from_list(docs)
+
+
+def load_t2_qa(subset: str) -> Dataset:
+    return Dataset.from_list(_load_rows(subset))
+
+
+def map_t2_document(doc: Mapping[str, Any]) -> DocumentPayload | None:
+    pdf_path = download_t2_pdf(doc["subset"], doc["split"], doc["file_name"])
+
+    metadata: dict[str, Any] = {"file_name": doc["file_name"]}
+    for key in ("company_name", "company_symbol", "report_year", "company_sector"):
+        value = doc.get(key)
+        if value is not None:
+            metadata[key] = value
+
+    title_parts = [
+        str(doc[key]) for key in ("company_name", "report_year") if doc.get(key)
+    ]
+    title = " ".join(title_parts) if title_parts else doc["context_id"]
+
+    return DocumentPayload(
+        uri=doc["context_id"],
+        source_path=pdf_path,
+        title=title,
+        metadata=metadata,
+    )
+
+
+def map_t2_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
+    return RetrievalSample(
+        question=doc["question"],
+        expected_uris=(doc["context_id"],),
+    )
+
+
+def build_t2_case(index: int, doc: Mapping[str, Any]) -> Case[str, str, dict[str, str]]:
+    metadata = {
+        "case_index": str(index),
+        "context_id": doc["context_id"],
+        "id": doc["id"],
+    }
+
+    return Case(
+        name=f"{index}_{doc['id']}",
+        inputs=doc["question"],
+        expected_output=str(doc["program_answer"]),
+        metadata=metadata,
+    )
+
+
+def _t2_spec(subset: str, key: str, db_filename: str) -> DatasetSpec:
+    return DatasetSpec(
+        key=key,
+        db_filename=db_filename,
+        document_loader=partial(load_t2_corpus, subset),
+        document_mapper=map_t2_document,
+        qa_loader=partial(load_t2_qa, subset),
+        qa_case_builder=build_t2_case,
+        retrieval_loader=partial(load_t2_qa, subset),
+        retrieval_mapper=map_t2_retrieval,
+        retrieval_evaluator=MAPEvaluator(),
+    )
+
+
+T2_FINQA_SPEC = _t2_spec(
+    subset="FinQA",
+    key="t2_finqa",
+    db_filename="t2_ragbench_finqa.lancedb",
+)

evaluations/tests/test_datasets.py

@@ -7,6 +7,13 @@ from evaluations.datasets.open_rag_bench import (
     map_orb_document,
     map_orb_retrieval,
 )
+from evaluations.datasets.t2_ragbench import (
+    build_t2_case,
+    download_t2_pdf,
+    load_t2_corpus,
+    map_t2_document,
+    map_t2_retrieval,
+)
 from evaluations.datasets.wix import (
     build_wix_case,
     map_wix_document,
@@ -158,3 +165,122 @@ class TestOpenRAGBench:
         assert is_multimodal_query("image") is True
         assert is_multimodal_query("image_table") is True
         assert is_multimodal_query("text") is False
+
+
+class TestT2RAGBench:
+    def _row(self, **overrides: object) -> dict[str, object]:
+        row: dict[str, object] = {
+            "id": "finqa_dev_0",
+            "context_id": "finqa_dev_ctx_138",
+            "subset": "FinQA",
+            "split": "dev",
+            "file_name": "pdf/V/2008/page_17.pdf",
+            "question": "What was the average payment volume per transaction?",
+            "program_answer": "127.4",
+            "company_name": "Visa Inc.",
+            "company_symbol": "V",
+            "report_year": 2008,
+            "company_sector": "Financials",
+        }
+        row.update(overrides)
+        return row
+
+    def test_map_document(self, tmp_path: Path) -> None:
+        pdf_path = tmp_path / "page_17.pdf"
+        pdf_path.write_bytes(b"%PDF-fake")
+        from unittest.mock import patch
+
+        with patch(
+            "evaluations.datasets.t2_ragbench.download_t2_pdf",
+            return_value=pdf_path,
+        ) as download:
+            payload = map_t2_document(self._row())
+
+        download.assert_called_once_with("FinQA", "dev", "pdf/V/2008/page_17.pdf")
+        assert payload is not None
+        assert payload.uri == "finqa_dev_ctx_138"
+        assert payload.source_path == pdf_path
+        assert payload.content is None
+        assert payload.title == "Visa Inc. 2008"
+        assert payload.metadata == {
+            "file_name": "pdf/V/2008/page_17.pdf",
+            "company_name": "Visa Inc.",
+            "company_symbol": "V",
+            "report_year": 2008,
+            "company_sector": "Financials",
+        }
+
+    def test_map_document_title_falls_back_to_context_id(self, tmp_path: Path) -> None:
+        pdf_path = tmp_path / "page.pdf"
+        pdf_path.write_bytes(b"%PDF-fake")
+        from unittest.mock import patch
+
+        row = self._row(company_name=None, report_year=None)
+        with patch(
+            "evaluations.datasets.t2_ragbench.download_t2_pdf",
+            return_value=pdf_path,
+        ):
+            payload = map_t2_document(row)
+
+        assert payload is not None
+        assert payload.title == "finqa_dev_ctx_138"
+        assert payload.metadata is not None
+        assert "company_name" not in payload.metadata
+        assert "report_year" not in payload.metadata
+
+    def test_map_retrieval(self) -> None:
+        sample = map_t2_retrieval(self._row())
+        assert sample is not None
+        assert sample.question == self._row()["question"]
+        assert sample.expected_uris == ("finqa_dev_ctx_138",)
+
+    def test_build_case(self) -> None:
+        case = build_t2_case(3, self._row())
+        assert case.name == "3_finqa_dev_0"
+        assert case.inputs == self._row()["question"]
+        assert case.expected_output == "127.4"
+        assert case.metadata is not None
+        assert case.metadata["case_index"] == "3"
+        assert case.metadata["context_id"] == "finqa_dev_ctx_138"
+
+    def test_build_case_casts_numeric_answer(self) -> None:
+        case = build_t2_case(0, self._row(program_answer=127.4))
+        assert case.expected_output == "127.4"
+
+    def test_download_pdf_materializes_pdf_suffix(self, tmp_path: Path) -> None:
+        # hf_hub_download can return a suffix-less blob path; the cached copy
+        # must carry the .pdf extension the converter dispatches on.
+        blob = tmp_path / "blobs" / "6aa49306deadbeef"
+        blob.parent.mkdir()
+        blob.write_bytes(b"%PDF-fake")
+        cache = tmp_path / "cache"
+        cache.mkdir()
+        from unittest.mock import patch
+
+        with (
+            patch("evaluations.datasets.t2_ragbench.get_cache_dir", return_value=cache),
+            patch(
+                "evaluations.datasets.t2_ragbench.hf_hub_download",
+                return_value=str(blob),
+            ),
+        ):
+            out = download_t2_pdf("FinQA", "dev", "pdf/V/2008/page_17.pdf")
+
+        assert out.suffix == ".pdf"
+        assert out.exists()
+        assert out.read_bytes() == b"%PDF-fake"
+        assert out.name == "FinQA_dev_pdf_V_2008_page_17.pdf"
+
+    def test_load_corpus_dedupes_by_context_id(self) -> None:
+        from unittest.mock import patch
+
+        rows = [
+            self._row(id="finqa_dev_0", context_id="ctx_a"),
+            self._row(id="finqa_dev_1", context_id="ctx_a"),
+            self._row(id="finqa_dev_2", context_id="ctx_b"),
+        ]
+        with patch("evaluations.datasets.t2_ragbench._load_rows", return_value=rows):
+            corpus = load_t2_corpus("FinQA")
+
+        assert len(corpus) == 2
+        assert {r["context_id"] for r in corpus} == {"ctx_a", "ctx_b"}