From da0ff433de20520ee68675e87d57f5eeb5d12040 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 24 May 2026 11:02:09 +0300 Subject: [PATCH] add MMLongBench-Doc evaluation dataset --- evaluations/evaluations/datasets/__init__.py | 2 + .../evaluations/datasets/mmlongbench.py | 153 ++++++++++++++++++ evaluations/tests/test_datasets.py | 122 ++++++++++++++ 3 files changed, 277 insertions(+) create mode 100644 evaluations/evaluations/datasets/mmlongbench.py diff --git a/evaluations/evaluations/datasets/__init__.py b/evaluations/evaluations/datasets/__init__.py index c973aa5f..9a9f9671 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/evaluations/evaluations/datasets/__init__.py @@ -1,5 +1,6 @@ from evaluations.config import DatasetSpec +from .mmlongbench import MMLONGBENCH_SPEC from .open_rag_bench import ORB_MULTIMODAL_SPEC, ORB_TEXT_SPEC from .wix import WIX_SPEC @@ -9,6 +10,7 @@ DATASETS: dict[str, DatasetSpec] = { WIX_SPEC, ORB_TEXT_SPEC, ORB_MULTIMODAL_SPEC, + MMLONGBENCH_SPEC, ) } diff --git a/evaluations/evaluations/datasets/mmlongbench.py b/evaluations/evaluations/datasets/mmlongbench.py new file mode 100644 index 00000000..120cba90 --- /dev/null +++ b/evaluations/evaluations/datasets/mmlongbench.py @@ -0,0 +1,153 @@ +import ast +import logging +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from datasets import Dataset, load_dataset +from huggingface_hub import snapshot_download +from pydantic_evals import Case + +from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample +from evaluations.evaluators import MAPEvaluator + +logger = logging.getLogger(__name__) + +REPO_ID = "yubo2333/MMLongBench-Doc" +PDF_SUBDIR = "documents" +_LIST_FIELDS = ("evidence_pages", "evidence_sources") + + +def get_cache_dir() -> Path: + cache_dir = Path.home() / ".cache" / "haiku.rag" / "evaluations" / "mmlongbench" + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + + +def ensure_pdfs_downloaded() -> Path: + cache_dir = get_cache_dir() + snapshot_download( + repo_id=REPO_ID, + repo_type="dataset", + allow_patterns=f"{PDF_SUBDIR}/*.pdf", + local_dir=cache_dir, + ) + return cache_dir / PDF_SUBDIR + + +def _parse_list_field(value: Any) -> list[Any]: + if isinstance(value, list): + return value + if not value: + return [] + return ast.literal_eval(value) + + +def _load_hf_qa_split() -> list[dict[str, Any]]: + dataset = load_dataset(REPO_ID, split="train") + return [dict(row) for row in dataset] + + +_qa_records: list[dict[str, Any]] | None = None + + +def load_qa_records() -> list[dict[str, Any]]: + global _qa_records + if _qa_records is not None: + return _qa_records + rows = _load_hf_qa_split() + for row in rows: + for field in _LIST_FIELDS: + row[field] = _parse_list_field(row.get(field)) + _qa_records = rows + return _qa_records + + +def load_mmlb_corpus() -> Dataset: + pdf_dir = ensure_pdfs_downloaded() + records = load_qa_records() + doc_types: dict[str, str] = {} + for row in records: + doc_id = row["doc_id"] + if doc_id not in doc_types: + doc_types[doc_id] = row.get("doc_type", "") + corpus = [ + {"doc_id": doc_id, "doc_type": doc_type} + for doc_id, doc_type in doc_types.items() + if (pdf_dir / doc_id).exists() + ] + return Dataset.from_list(corpus) + + +def map_mmlb_document(doc: Mapping[str, Any]) -> DocumentPayload | None: + doc_id = doc["doc_id"] + pdf_path = get_cache_dir() / PDF_SUBDIR / doc_id + if not pdf_path.exists(): + logger.warning(f"PDF not found in cache: {doc_id}") + return None + return DocumentPayload( + uri=doc_id, + source_path=pdf_path, + title=doc_id, + metadata={"doc_type": doc.get("doc_type", "")}, + ) + + +def load_mmlb_qa() -> Dataset: + return Dataset.from_list(load_qa_records()) + + +def build_mmlb_case( + index: int, doc: Mapping[str, Any] +) -> Case[str, str, dict[str, str]]: + evidence_pages = list(doc.get("evidence_pages") or []) + evidence_sources = list(doc.get("evidence_sources") or []) + metadata: dict[str, str] = { + "case_index": str(index), + "doc_id": doc["doc_id"], + "doc_type": doc.get("doc_type", ""), + "answer_format": doc.get("answer_format", ""), + "evidence_pages": str(evidence_pages), + "evidence_sources": ",".join(str(s) for s in evidence_sources), + } + return Case( + name=f"{index}_{doc['doc_id']}", + inputs=doc["question"], + expected_output=doc["answer"], + metadata=metadata, + ) + + +def load_mmlb_retrieval() -> Dataset: + records = [] + for row in load_qa_records(): + if not row.get("evidence_pages"): + continue + records.append(row) + return Dataset.from_list(records) + + +def map_mmlb_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None: + evidence_pages = doc.get("evidence_pages") or [] + if not evidence_pages: + return None + sources = doc.get("evidence_sources") or [] + source_type = ",".join(str(s) for s in sources) if sources else None + return RetrievalSample( + question=doc["question"], + expected_uris=(doc["doc_id"],), + source_type=source_type, + ) + + +MMLONGBENCH_SPEC = DatasetSpec( + key="mmlongbench", + db_filename="mmlongbench.lancedb", + document_loader=load_mmlb_corpus, + document_mapper=map_mmlb_document, + qa_loader=load_mmlb_qa, + qa_case_builder=build_mmlb_case, + retrieval_loader=load_mmlb_retrieval, + retrieval_mapper=map_mmlb_retrieval, + retrieval_evaluator=MAPEvaluator(), +) diff --git a/evaluations/tests/test_datasets.py b/evaluations/tests/test_datasets.py index 07a516ee..82e792c8 100644 --- a/evaluations/tests/test_datasets.py +++ b/evaluations/tests/test_datasets.py @@ -1,5 +1,11 @@ from pathlib import Path +from evaluations.datasets.mmlongbench import ( + build_mmlb_case, + load_qa_records, + map_mmlb_document, + map_mmlb_retrieval, +) from evaluations.datasets.open_rag_bench import ( build_orb_case, download_pdf, @@ -158,3 +164,119 @@ 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 TestMMLongBenchDoc: + def test_map_document(self, tmp_path: Path) -> None: + pdf_dir = tmp_path / "documents" + pdf_dir.mkdir() + pdf_path = pdf_dir / "report.pdf" + pdf_path.write_bytes(b"%PDF-fake") + + from unittest.mock import patch + + with patch( + "evaluations.datasets.mmlongbench.get_cache_dir", + return_value=tmp_path, + ): + payload = map_mmlb_document( + {"doc_id": "report.pdf", "doc_type": "Financial report"} + ) + + assert payload is not None + assert payload.uri == "report.pdf" + assert payload.title == "report.pdf" + assert payload.source_path == pdf_path + assert payload.metadata == {"doc_type": "Financial report"} + + def test_map_document_missing_pdf(self, tmp_path: Path) -> None: + from unittest.mock import patch + + with patch( + "evaluations.datasets.mmlongbench.get_cache_dir", + return_value=tmp_path, + ): + payload = map_mmlb_document( + {"doc_id": "missing.pdf", "doc_type": "Brochure"} + ) + + assert payload is None + + def test_map_retrieval(self) -> None: + doc = { + "question": "What is the revenue?", + "doc_id": "NIKE_2021_10K.pdf", + "evidence_pages": [3, 5], + "evidence_sources": ["Table", "Pure-text"], + } + sample = map_mmlb_retrieval(doc) + assert sample is not None + assert sample.question == "What is the revenue?" + assert sample.expected_uris == ("NIKE_2021_10K.pdf",) + assert sample.source_type == "Table,Pure-text" + + def test_map_retrieval_skips_unanswerable(self) -> None: + doc = { + "question": "What does the document say about Mars?", + "doc_id": "NIKE_2021_10K.pdf", + "evidence_pages": [], + "evidence_sources": [], + } + assert map_mmlb_retrieval(doc) is None + + def test_build_case(self) -> None: + doc = { + "doc_id": "report.pdf", + "doc_type": "Financial report", + "question": "What is the net income?", + "answer": "42", + "evidence_pages": [5], + "evidence_sources": ["Table"], + "answer_format": "Int", + } + case = build_mmlb_case(7, doc) + assert case.name == "7_report.pdf" + assert case.inputs == "What is the net income?" + assert case.expected_output == "42" + assert case.metadata is not None + assert case.metadata["doc_id"] == "report.pdf" + assert case.metadata["doc_type"] == "Financial report" + assert case.metadata["answer_format"] == "Int" + assert case.metadata["evidence_pages"] == "[5]" + assert case.metadata["evidence_sources"] == "Table" + assert case.metadata["case_index"] == "7" + + def test_load_qa_records_parses_list_fields(self) -> None: + from unittest.mock import patch + + raw_rows = [ + { + "doc_id": "a.pdf", + "doc_type": "Brochure", + "question": "Q1?", + "answer": "A1", + "evidence_pages": "[3, 5]", + "evidence_sources": "['Table', 'Pure-text']", + "answer_format": "Str", + }, + { + "doc_id": "b.pdf", + "doc_type": "Academic paper", + "question": "Q2?", + "answer": "Not answerable", + "evidence_pages": "[]", + "evidence_sources": "[]", + "answer_format": "None", + }, + ] + + with patch( + "evaluations.datasets.mmlongbench._load_hf_qa_split", + return_value=raw_rows, + ): + records = load_qa_records() + + assert records[0]["evidence_pages"] == [3, 5] + assert records[0]["evidence_sources"] == ["Table", "Pure-text"] + assert records[1]["evidence_pages"] == [] + assert records[1]["evidence_sources"] == []