58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
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
|
|
|
|
|
|
def load_repliqa_corpus() -> Dataset:
|
|
dataset_dict = cast(DatasetDict, load_dataset("ServiceNow/repliqa"))
|
|
dataset = cast(Dataset, dataset_dict["repliqa_3"])
|
|
return dataset.filter(lambda doc: doc["document_topic"] == "News Stories")
|
|
|
|
|
|
def map_repliqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
|
return DocumentPayload(
|
|
uri=str(doc["document_id"]),
|
|
content=doc["document_extracted"],
|
|
)
|
|
|
|
|
|
def map_repliqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
|
expected_answer = doc["answer"]
|
|
if expected_answer == "The answer is not found in the document.":
|
|
return None
|
|
return RetrievalSample(
|
|
question=doc["question"],
|
|
expected_uris=(str(doc["document_id"]),),
|
|
)
|
|
|
|
|
|
def build_repliqa_case(
|
|
index: int, doc: Mapping[str, Any]
|
|
) -> Case[str, str, dict[str, str]]:
|
|
document_id = doc["document_id"]
|
|
case_name = f"{index}_{document_id}" if document_id is not None else f"case_{index}"
|
|
return Case(
|
|
name=case_name,
|
|
inputs=doc["question"],
|
|
expected_output=doc["answer"],
|
|
metadata={
|
|
"document_id": str(document_id),
|
|
"case_index": str(index),
|
|
},
|
|
)
|
|
|
|
|
|
REPLIQ_SPEC = DatasetSpec(
|
|
key="repliqa",
|
|
db_filename="repliqa.lancedb",
|
|
document_loader=load_repliqa_corpus,
|
|
document_mapper=map_repliqa_document,
|
|
qa_loader=load_repliqa_corpus,
|
|
qa_case_builder=build_repliqa_case,
|
|
retrieval_loader=load_repliqa_corpus,
|
|
retrieval_mapper=map_repliqa_retrieval,
|
|
)
|