25 lines
729 B
Python
25 lines
729 B
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
from datasets import Dataset, load_dataset, load_from_disk
|
|
|
|
from .llm_judge import LLMJudge
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def qa_corpus() -> Dataset:
|
|
ds_path = Path(__file__).parent / "data" / "dataset"
|
|
ds_path.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
ds: Dataset = load_from_disk(ds_path) # type: ignore
|
|
return ds
|
|
except FileNotFoundError:
|
|
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore
|
|
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
|
|
corpus.save_to_disk(ds_path)
|
|
return corpus
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def llm_judge() -> LLMJudge:
|
|
return LLMJudge()
|