add MMLongBench-Doc evaluation dataset

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,
     )
 }
 

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(),
+)

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"] == []