81 lines
2.2 KiB
Python
81 lines
2.2 KiB
Python
import json
|
|
from collections.abc import Iterable, 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_wix_corpus() -> Dataset:
|
|
dataset_dict = cast(DatasetDict, load_dataset("Wix/WixQA", "wix_kb_corpus"))
|
|
return cast(Dataset, dataset_dict["train"])
|
|
|
|
|
|
def map_wix_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
|
article_id = doc.get("id")
|
|
url = doc.get("url")
|
|
uri = str(article_id) if article_id is not None else str(url)
|
|
|
|
metadata: dict[str, str] = {}
|
|
if article_id is not None:
|
|
metadata["article_id"] = str(article_id)
|
|
if url:
|
|
metadata["url"] = str(url)
|
|
|
|
return DocumentPayload(
|
|
uri=uri,
|
|
content=doc["contents"],
|
|
title=doc.get("title"),
|
|
metadata=metadata or None,
|
|
)
|
|
|
|
|
|
def load_wix_qa() -> Dataset:
|
|
dataset_dict = cast(DatasetDict, load_dataset("Wix/WixQA", "wixqa_expertwritten"))
|
|
return cast(Dataset, dataset_dict["train"])
|
|
|
|
|
|
def map_wix_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
|
article_ids: Iterable[int | str] | None = doc.get("article_ids")
|
|
if not article_ids:
|
|
return None
|
|
|
|
expected_uris = tuple(str(article_id) for article_id in article_ids)
|
|
return RetrievalSample(
|
|
question=doc["question"],
|
|
expected_uris=expected_uris,
|
|
)
|
|
|
|
|
|
def build_wix_case(
|
|
index: int, doc: Mapping[str, Any]
|
|
) -> Case[str, str, dict[str, str]]:
|
|
article_ids = tuple(str(article_id) for article_id in doc.get("article_ids") or [])
|
|
joined_ids = "-".join(article_ids)
|
|
case_name = f"{index}_{joined_ids}" if joined_ids else f"case_{index}"
|
|
|
|
metadata = {
|
|
"case_index": str(index),
|
|
"document_ids": json.dumps(article_ids),
|
|
}
|
|
|
|
return Case(
|
|
name=case_name,
|
|
inputs=doc["question"],
|
|
expected_output=doc["answer"],
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
WIX_SPEC = DatasetSpec(
|
|
key="wix",
|
|
db_filename="wix.lancedb",
|
|
document_loader=load_wix_corpus,
|
|
document_mapper=map_wix_document,
|
|
qa_loader=load_wix_qa,
|
|
qa_case_builder=build_wix_case,
|
|
retrieval_loader=load_wix_qa,
|
|
retrieval_mapper=map_wix_retrieval,
|
|
)
|