From 752673505979e0e6260d4e0bc0bd73242f1682ec Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Dec 2025 10:18:44 +0200 Subject: [PATCH] hotpotqa adapter --- evaluations/evaluations/datasets/__init__.py | 5 +- evaluations/evaluations/datasets/hotpotqa.py | 108 +++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) create mode 100644 evaluations/evaluations/datasets/hotpotqa.py diff --git a/evaluations/evaluations/datasets/__init__.py b/evaluations/evaluations/datasets/__init__.py index 94a00c34..54b92b3b 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/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"] diff --git a/evaluations/evaluations/datasets/hotpotqa.py b/evaluations/evaluations/datasets/hotpotqa.py new file mode 100644 index 00000000..2f0a9b1c --- /dev/null +++ b/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(), +)