hotpotqa adapter

evaluations/evaluations/datasets/__init__.py

@@ -1,8 +1,11 @@
 from evaluations.config import DatasetSpec
 
+from .hotpotqa import HOTPOTQA_SPEC
 from .repliqa import REPLIQ_SPEC
 from .wix import WIX_SPEC
 
-DATASETS: dict[str, DatasetSpec] = {spec.key: spec for spec in (REPLIQ_SPEC, WIX_SPEC)}
+DATASETS: dict[str, DatasetSpec] = {
+    spec.key: spec for spec in (REPLIQ_SPEC, WIX_SPEC, HOTPOTQA_SPEC)
+}
 
 __all__ = ["DATASETS"]

evaluations/evaluations/datasets/hotpotqa.py

@@ -0,0 +1,108 @@
+from collections.abc import Mapping
+from typing import Any, cast
+
+from datasets import Dataset, DatasetDict, load_dataset
+from pydantic_evals import Case
+
+from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
+from evaluations.evaluators import MAPEvaluator
+
+
+def load_hotpotqa_validation() -> Dataset:
+    dataset_dict = cast(DatasetDict, load_dataset("hotpotqa/hotpot_qa", "distractor"))
+    return cast(Dataset, dataset_dict["validation"])
+
+
+def extract_unique_documents(dataset: Dataset) -> list[dict[str, Any]]:
+    """Extract unique documents from all context paragraphs, deduplicated by title."""
+    seen_titles: set[str] = set()
+    documents: list[dict[str, Any]] = []
+
+    for sample in dataset:
+        sample = cast(Mapping[str, Any], sample)
+        context = sample["context"]
+        titles = context["title"]
+        sentences_list = context["sentences"]
+
+        for title, sentences in zip(titles, sentences_list):
+            if title in seen_titles:
+                continue
+            seen_titles.add(title)
+            content = " ".join(sentences)
+            documents.append({"title": title, "content": content})
+
+    return documents
+
+
+_cached_documents: list[dict[str, Any]] | None = None
+
+
+def load_hotpotqa_documents() -> list[dict[str, Any]]:
+    """Load and cache unique documents from HotpotQA."""
+    global _cached_documents
+    if _cached_documents is None:
+        dataset = load_hotpotqa_validation()
+        _cached_documents = extract_unique_documents(dataset)
+    return _cached_documents
+
+
+def document_loader() -> Dataset:
+    """Return documents as a Dataset-like iterable."""
+    docs = load_hotpotqa_documents()
+    return Dataset.from_list(docs)
+
+
+def map_hotpotqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
+    return DocumentPayload(
+        uri=doc["title"],
+        content=doc["content"],
+        title=doc["title"],
+    )
+
+
+def map_hotpotqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
+    supporting_facts = doc["supporting_facts"]
+    titles = supporting_facts["title"]
+    if not titles:
+        return None
+
+    unique_titles = tuple(dict.fromkeys(titles))
+    return RetrievalSample(
+        question=doc["question"],
+        expected_uris=unique_titles,
+    )
+
+
+def build_hotpotqa_case(
+    index: int, doc: Mapping[str, Any]
+) -> Case[str, str, dict[str, str]]:
+    question_id = doc["id"]
+    question_type = doc["type"]
+    level = doc["level"]
+
+    case_name = f"{index}_{question_id}"
+
+    return Case(
+        name=case_name,
+        inputs=doc["question"],
+        expected_output=doc["answer"],
+        metadata={
+            "question_id": str(question_id),
+            "type": str(question_type),
+            "level": str(level),
+            "case_index": str(index),
+        },
+    )
+
+
+HOTPOTQA_SPEC = DatasetSpec(
+    key="hotpotqa",
+    db_filename="hotpotqa.lancedb",
+    document_loader=document_loader,
+    document_mapper=map_hotpotqa_document,
+    qa_loader=load_hotpotqa_validation,
+    qa_case_builder=build_hotpotqa_case,
+    retrieval_loader=load_hotpotqa_validation,
+    retrieval_mapper=map_hotpotqa_retrieval,
+    retrieval_evaluator=MAPEvaluator(),
+)