diff --git a/CHANGELOG.md b/CHANGELOG.md index 4490b25e..d1276c01 100644 --- a/CHANGELOG.md +++ b/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`. diff --git a/evaluations/evaluations/datasets/__init__.py b/evaluations/evaluations/datasets/__init__.py index c973aa5f..55cc8604 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/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, ) } diff --git a/evaluations/evaluations/datasets/t2_ragbench.py b/evaluations/evaluations/datasets/t2_ragbench.py new file mode 100644 index 00000000..21ee0adf --- /dev/null +++ b/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", +) diff --git a/evaluations/tests/test_datasets.py b/evaluations/tests/test_datasets.py index 07a516ee..df1c8344 100644 --- a/evaluations/tests/test_datasets.py +++ b/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"}