Merge pull request #87 from ggozad/feat/multiple-eval-datasets
Run evaluations with alternate datasets. Use WixQA as an example.
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
2ae11200b4
17 changed files with 562 additions and 270 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -11,7 +11,7 @@ wheels/
|
|||
|
||||
# tests
|
||||
.coverage*
|
||||
tests/data/
|
||||
evaluations/data/
|
||||
.pytest_cache/
|
||||
.ruff_cache/
|
||||
|
||||
|
|
|
|||
|
|
@ -2,8 +2,9 @@
|
|||
|
||||
We use the [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) dataset for the evaluation of `haiku.rag`.
|
||||
|
||||
You can perform your own evaluations using as example the script found at
|
||||
`tests/generate_benchmark_db.py`. The evaluation flow is orchestrated with
|
||||
You can perform your own evaluations with the Typer CLI in
|
||||
`evaluations/benchmark.py`, for example `python -m evaluations.benchmark repliqa`.
|
||||
The evaluation flow is orchestrated with
|
||||
[`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals),
|
||||
which we leverage for dataset management, scoring, and report generation.
|
||||
|
||||
|
|
@ -16,13 +17,11 @@ The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 fo
|
|||
|
||||
| Embedding Model | Document in top 1 | Document in top 3 | Reranker |
|
||||
|---------------------------------------|-------------------|-------------------|------------------------|
|
||||
| Ollama / `qwen3-embedding` | 0.81 | 0.95 | None |
|
||||
| Ollama / `qwen3-embedding` | 0.91 | 0.98 | `mxbai-rerank-base-v2` |
|
||||
| Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None |
|
||||
| Ollama / `mxbai-embed-large` | 0.90 | 0.95 | `mxbai-rerank-base-v2` |
|
||||
| Ollama / `nomic-embed-text-v1.5` | 0.74 | 0.90 | None |
|
||||
| Ollama / `qwen3-embedding` | 0.81 | 0.95 | None |
|
||||
<!-- | OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
|
||||
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
|
||||
| OpenAI / `text-embeddings-3-small` | 0.83 | 0.90 | Cohere / `rerank-v3.5` | -->
|
||||
|
||||
## Question/Answer evaluation
|
||||
|
||||
|
|
@ -38,5 +37,29 @@ determine whether the answer is correct. The obtained accuracy is as follows:
|
|||
| Ollama / `mxbai-embed-large` | Ollama / `qwen3:0.6b` | 0.28 | None |
|
||||
|
||||
Note the significant degradation when very small models are used such as `qwen3:0.6b`.
|
||||
<!-- | Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None |
|
||||
| OpenAI / `text-embeddings-3-small` | OpenAI / `gpt-4-turbo` | 0.62 | None | -->
|
||||
|
||||
## Wix dataset
|
||||
|
||||
We also track retrieval performance on [WixQA](https://huggingface.co/datasets/Wix/WixQA),
|
||||
a dataset of real customer support questions paired with curated answers from
|
||||
Wix. The benchmark follows the evaluation protocol described in the
|
||||
[WixQA paper](https://arxiv.org/abs/2505.08643) and gives us a view into how the
|
||||
system handles conversational, product-specific support queries.
|
||||
|
||||
For recall, we index the reference answer passages shipped with the dataset and
|
||||
run retrieval against each user question. Each sample supplies one or more
|
||||
relevant passage URIs; we count how many of those URIs land inside the top *k*
|
||||
retrieved documents, divide by the number of relevant passages for that query,
|
||||
and average across all queries.
|
||||
|
||||
The results for recall using the `WixQA` dataset are as follows:
|
||||
|
||||
| Embedding Model | Document in top 1 | Document in top 3 | Reranker |
|
||||
|----------------------------|-------------------|-------------------|------------------------|
|
||||
| `qwen3-embedding` | 0.36 | 0.57 | `mxbai-rerank-base-v2` |
|
||||
|
||||
And for QA accuracy,
|
||||
|
||||
| Embedding Model | QA Model | Accuracy | Reranker |
|
||||
|----------------------------|-----------|----------|------------------------|
|
||||
| `qwen3-embedding` | `gpt-oss` | 0.75 | `mxbai-rerank-base-v2` |
|
||||
|
|
|
|||
0
evaluations/__init__.py
Normal file
0
evaluations/__init__.py
Normal file
320
evaluations/benchmark.py
Normal file
320
evaluations/benchmark.py
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, cast
|
||||
|
||||
import logfire
|
||||
import typer
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
from pydantic_evals import Dataset as EvalDataset
|
||||
from pydantic_evals.evaluators import IsInstance, LLMJudge
|
||||
from pydantic_evals.reporting import ReportCaseFailure
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
|
||||
from evaluations.config import DatasetSpec, RetrievalSample
|
||||
from evaluations.datasets import DATASETS
|
||||
from evaluations.llm_judge import ANSWER_EQUIVALENCE_RUBRIC
|
||||
from haiku.rag import logging # noqa: F401
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.logging import configure_cli_logging
|
||||
from haiku.rag.qa import get_qa_agent
|
||||
|
||||
QA_JUDGE_MODEL = "qwen3"
|
||||
|
||||
logfire.configure(send_to_logfire="if-token-present", service_name="evals")
|
||||
logfire.instrument_pydantic_ai()
|
||||
configure_cli_logging()
|
||||
console = Console()
|
||||
|
||||
|
||||
async def populate_db(spec: DatasetSpec) -> None:
|
||||
spec.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
corpus = spec.document_loader()
|
||||
if spec.document_limit is not None:
|
||||
corpus = corpus.select(range(min(spec.document_limit, len(corpus))))
|
||||
|
||||
with Progress() as progress:
|
||||
task = progress.add_task("[green]Populating database...", total=len(corpus))
|
||||
async with HaikuRAG(spec.db_path) as rag:
|
||||
for doc in corpus:
|
||||
doc_mapping = cast(Mapping[str, Any], doc)
|
||||
payload = spec.document_mapper(doc_mapping)
|
||||
if payload is None:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
existing = await rag.get_document_by_uri(payload.uri)
|
||||
if existing is not None:
|
||||
assert existing.id
|
||||
chunks = await rag.chunk_repository.get_by_document_id(existing.id)
|
||||
if chunks:
|
||||
progress.advance(task)
|
||||
continue
|
||||
await rag.document_repository.delete(existing.id)
|
||||
|
||||
await rag.create_document(
|
||||
content=payload.content,
|
||||
uri=payload.uri,
|
||||
title=payload.title,
|
||||
metadata=payload.metadata,
|
||||
)
|
||||
progress.advance(task)
|
||||
rag.store.vacuum()
|
||||
|
||||
|
||||
def _is_relevant_match(retrieved_uri: str | None, sample: RetrievalSample) -> bool:
|
||||
return retrieved_uri is not None and retrieved_uri in sample.expected_uris
|
||||
|
||||
|
||||
async def run_retrieval_benchmark(spec: DatasetSpec) -> dict[str, float] | None:
|
||||
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
|
||||
console.print("Skipping retrieval benchmark; no retrieval config.")
|
||||
return None
|
||||
|
||||
corpus = spec.retrieval_loader()
|
||||
|
||||
recall_totals = {
|
||||
1: 0.0,
|
||||
3: 0.0,
|
||||
5: 0.0,
|
||||
}
|
||||
total_queries = 0
|
||||
|
||||
with Progress() as progress:
|
||||
task = progress.add_task(
|
||||
"[blue]Running retrieval benchmark...", total=len(corpus)
|
||||
)
|
||||
async with HaikuRAG(spec.db_path) as rag:
|
||||
for doc in corpus:
|
||||
doc_mapping = cast(Mapping[str, Any], doc)
|
||||
sample = spec.retrieval_mapper(doc_mapping)
|
||||
if sample is None or sample.skip:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
matches = await rag.search(query=sample.question, limit=5)
|
||||
if not matches:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
total_queries += 1
|
||||
|
||||
retrieved_uris: list[str] = []
|
||||
for chunk, _ in matches:
|
||||
if chunk.document_id is None:
|
||||
continue
|
||||
retrieved_doc = await rag.get_document_by_id(chunk.document_id)
|
||||
if retrieved_doc and retrieved_doc.uri:
|
||||
retrieved_uris.append(retrieved_doc.uri)
|
||||
|
||||
# Compute per-query recall@K by counting how many relevant
|
||||
# documents are retrieved within the first K results and
|
||||
# averaging these fractions across all queries.
|
||||
for cutoff in (1, 3, 5):
|
||||
top_k = set(retrieved_uris[:cutoff])
|
||||
relevant = set(sample.expected_uris)
|
||||
if relevant:
|
||||
matched = len(top_k & relevant)
|
||||
recall_totals[cutoff] += matched / len(relevant)
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
if total_queries == 0:
|
||||
console.print("No retrieval cases to evaluate.")
|
||||
return None
|
||||
|
||||
recall_at_1 = recall_totals[1] / total_queries
|
||||
recall_at_3 = recall_totals[3] / total_queries
|
||||
recall_at_5 = recall_totals[5] / total_queries
|
||||
|
||||
console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Total queries: {total_queries}")
|
||||
console.print(f"Recall@1: {recall_at_1:.4f}")
|
||||
console.print(f"Recall@3: {recall_at_3:.4f}")
|
||||
console.print(f"Recall@5: {recall_at_5:.4f}")
|
||||
|
||||
return {
|
||||
"recall@1": recall_at_1,
|
||||
"recall@3": recall_at_3,
|
||||
"recall@5": recall_at_5,
|
||||
}
|
||||
|
||||
|
||||
async def run_qa_benchmark(
|
||||
spec: DatasetSpec, qa_limit: int | None = None
|
||||
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
|
||||
corpus = spec.qa_loader()
|
||||
if qa_limit is not None:
|
||||
corpus = corpus.select(range(min(qa_limit, len(corpus))))
|
||||
|
||||
cases = [
|
||||
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
|
||||
for index, doc in enumerate(corpus, start=1)
|
||||
]
|
||||
|
||||
judge_model = OpenAIChatModel(
|
||||
model_name=QA_JUDGE_MODEL,
|
||||
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
|
||||
)
|
||||
|
||||
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
|
||||
cases=cases,
|
||||
evaluators=[
|
||||
IsInstance(type_name="str"),
|
||||
LLMJudge(
|
||||
rubric=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model=judge_model,
|
||||
assertion={
|
||||
"evaluation_name": "answer_equivalent",
|
||||
"include_reason": True,
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
total_processed = 0
|
||||
passing_cases = 0
|
||||
failures: list[ReportCaseFailure[str, str, dict[str, str]]] = []
|
||||
|
||||
with Progress(console=console) as progress:
|
||||
qa_task = progress.add_task(
|
||||
"[yellow]Evaluating QA cases...",
|
||||
total=len(evaluation_dataset.cases),
|
||||
)
|
||||
|
||||
async with HaikuRAG(spec.db_path) as rag:
|
||||
qa = get_qa_agent(rag)
|
||||
|
||||
async def answer_question(question: str) -> str:
|
||||
return await qa.answer(question)
|
||||
|
||||
for case in evaluation_dataset.cases:
|
||||
progress.console.print(f"\n[bold]Evaluating case:[/bold] {case.name}")
|
||||
|
||||
single_case_dataset = EvalDataset[str, str, dict[str, str]](
|
||||
cases=[case],
|
||||
evaluators=evaluation_dataset.evaluators,
|
||||
)
|
||||
|
||||
report = await single_case_dataset.evaluate(
|
||||
answer_question,
|
||||
name="qa_answer",
|
||||
max_concurrency=1,
|
||||
progress=False,
|
||||
)
|
||||
|
||||
total_processed += 1
|
||||
|
||||
if report.cases:
|
||||
result_case = report.cases[0]
|
||||
|
||||
equivalence = result_case.assertions.get("answer_equivalent")
|
||||
progress.console.print(f"Question: {result_case.inputs}")
|
||||
progress.console.print(f"Expected: {result_case.expected_output}")
|
||||
progress.console.print(f"Generated: {result_case.output}")
|
||||
if equivalence is not None:
|
||||
progress.console.print(
|
||||
f"Equivalent: {equivalence.value}"
|
||||
+ (f" — {equivalence.reason}" if equivalence.reason else "")
|
||||
)
|
||||
if equivalence.value:
|
||||
passing_cases += 1
|
||||
|
||||
progress.console.print("")
|
||||
|
||||
if report.failures:
|
||||
failures.extend(report.failures)
|
||||
failure = report.failures[0]
|
||||
progress.console.print(
|
||||
"[red]Failure encountered during case evaluation:[/red]"
|
||||
)
|
||||
progress.console.print(f"Question: {failure.inputs}")
|
||||
progress.console.print(f"Error: {failure.error_message}")
|
||||
progress.console.print("")
|
||||
|
||||
progress.console.print(
|
||||
f"[green]Accuracy: {(passing_cases / total_processed):.4f} "
|
||||
f"{passing_cases}/{total_processed}[/green]"
|
||||
)
|
||||
progress.advance(qa_task)
|
||||
|
||||
total_cases = total_processed
|
||||
accuracy = passing_cases / total_cases if total_cases > 0 else 0
|
||||
|
||||
console.print("\n=== QA Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Total questions: {total_cases}")
|
||||
console.print(f"Correct answers: {passing_cases}")
|
||||
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
|
||||
|
||||
if failures:
|
||||
console.print("[red]\nSummary of failures:[/red]")
|
||||
for failure in failures:
|
||||
console.print(f"Case: {failure.name}")
|
||||
console.print(f"Question: {failure.inputs}")
|
||||
console.print(f"Error: {failure.error_message}")
|
||||
console.print("")
|
||||
|
||||
return failures[0] if failures else None
|
||||
|
||||
|
||||
async def evaluate_dataset(
|
||||
spec: DatasetSpec,
|
||||
skip_db: bool,
|
||||
skip_retrieval: bool,
|
||||
skip_qa: bool,
|
||||
qa_limit: int | None,
|
||||
) -> None:
|
||||
if not skip_db:
|
||||
console.print(f"Using dataset: {spec.key}", style="bold magenta")
|
||||
await populate_db(spec)
|
||||
|
||||
if not skip_retrieval:
|
||||
console.print("Running retrieval benchmarks...", style="bold blue")
|
||||
await run_retrieval_benchmark(spec)
|
||||
|
||||
if not skip_qa:
|
||||
console.print("\nRunning QA benchmarks...", style="bold yellow")
|
||||
await run_qa_benchmark(spec, qa_limit=qa_limit)
|
||||
|
||||
|
||||
app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.")
|
||||
|
||||
|
||||
@app.command()
|
||||
def run(
|
||||
dataset: str = typer.Argument(..., help="Dataset key to evaluate."),
|
||||
skip_db: bool = typer.Option(
|
||||
False, "--skip-db", help="Skip updateing the evaluation db."
|
||||
),
|
||||
skip_retrieval: bool = typer.Option(
|
||||
False, "--skip-retrieval", help="Skip retrieval benchmark."
|
||||
),
|
||||
skip_qa: bool = typer.Option(False, "--skip-qa", help="Skip QA benchmark."),
|
||||
qa_limit: int | None = typer.Option(
|
||||
None, "--qa-limit", help="Limit number of QA cases."
|
||||
),
|
||||
) -> None:
|
||||
spec = DATASETS.get(dataset.lower())
|
||||
if spec is None:
|
||||
valid_datasets = ", ".join(sorted(DATASETS))
|
||||
raise typer.BadParameter(
|
||||
f"Unknown dataset '{dataset}'. Choose from: {valid_datasets}"
|
||||
)
|
||||
|
||||
asyncio.run(
|
||||
evaluate_dataset(
|
||||
spec=spec,
|
||||
skip_db=skip_db,
|
||||
skip_retrieval=skip_retrieval,
|
||||
skip_qa=skip_qa,
|
||||
qa_limit=qa_limit,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
46
evaluations/config.py
Normal file
46
evaluations/config.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from datasets import Dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentPayload:
|
||||
uri: str
|
||||
content: str
|
||||
title: str | None = None
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetrievalSample:
|
||||
question: str
|
||||
expected_uris: tuple[str, ...]
|
||||
skip: bool = False
|
||||
|
||||
|
||||
DocumentLoader = Callable[[], Dataset]
|
||||
DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None]
|
||||
RetrievalLoader = Callable[[], Dataset]
|
||||
RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None]
|
||||
CaseBuilder = Callable[[int, Mapping[str, Any]], Case[str, str, dict[str, str]]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatasetSpec:
|
||||
key: str
|
||||
db_filename: str
|
||||
document_loader: DocumentLoader
|
||||
document_mapper: DocumentMapper
|
||||
qa_loader: DocumentLoader
|
||||
qa_case_builder: CaseBuilder
|
||||
retrieval_loader: RetrievalLoader | None = None
|
||||
retrieval_mapper: RetrievalMapper | None = None
|
||||
document_limit: int | None = None
|
||||
|
||||
@property
|
||||
def db_path(self) -> Path:
|
||||
return Path(__file__).parent / "data" / self.db_filename
|
||||
8
evaluations/datasets/__init__.py
Normal file
8
evaluations/datasets/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
from evaluations.config import DatasetSpec
|
||||
|
||||
from .repliqa import REPLIQ_SPEC
|
||||
from .wix import WIX_SPEC
|
||||
|
||||
DATASETS: dict[str, DatasetSpec] = {spec.key: spec for spec in (REPLIQ_SPEC, WIX_SPEC)}
|
||||
|
||||
__all__ = ["DATASETS"]
|
||||
58
evaluations/datasets/repliqa.py
Normal file
58
evaluations/datasets/repliqa.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
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,
|
||||
)
|
||||
81
evaluations/datasets/wix.py
Normal file
81
evaluations/datasets/wix.py
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
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,
|
||||
)
|
||||
|
|
@ -37,7 +37,7 @@ class LLMJudgeResponseSchema(BaseModel):
|
|||
class LLMJudge:
|
||||
"""LLM-as-judge for evaluating answer equivalence using Pydantic AI."""
|
||||
|
||||
def __init__(self, model: str = "qwen3"):
|
||||
def __init__(self, model: str = "gpt-oss"):
|
||||
# Create Ollama model
|
||||
ollama_model = OpenAIChatModel(
|
||||
model_name=model,
|
||||
|
|
@ -49,6 +49,7 @@ class LLMJudge:
|
|||
model=ollama_model,
|
||||
output_type=LLMJudgeResponseSchema,
|
||||
system_prompt=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
retries=3,
|
||||
)
|
||||
|
||||
async def judge_answers(
|
||||
|
|
@ -48,6 +48,9 @@ haiku-rag = "haiku.rag.cli:cli"
|
|||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build]
|
||||
exclude = ["/docs", "/evaluations", "/tests", "/.github"]
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/haiku"]
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ from haiku.rag.embeddings.base import EmbedderBase
|
|||
class Embedder(EmbedderBase):
|
||||
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
||||
client = AsyncOpenAI(base_url=f"{Config.OLLAMA_BASE_URL}/v1", api_key="dummy")
|
||||
if not text:
|
||||
return []
|
||||
response = await client.embeddings.create(
|
||||
model=self._model,
|
||||
input=text,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ from haiku.rag.embeddings.base import EmbedderBase
|
|||
class Embedder(EmbedderBase):
|
||||
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
||||
client = AsyncOpenAI()
|
||||
if not text:
|
||||
return []
|
||||
response = await client.embeddings.create(
|
||||
model=self._model,
|
||||
input=text,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ class Embedder(EmbedderBase):
|
|||
client = AsyncOpenAI(
|
||||
base_url=f"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy"
|
||||
)
|
||||
if not text:
|
||||
return []
|
||||
response = await client.embeddings.create(
|
||||
model=self._model,
|
||||
input=text,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ try:
|
|||
class Embedder(EmbedderBase):
|
||||
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
||||
client = Client()
|
||||
if not text:
|
||||
return []
|
||||
if isinstance(text, str):
|
||||
res = client.embed([text], model=self._model, output_dtype="float")
|
||||
return res.embeddings[0] # type: ignore[return-value]
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import os
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.reranking.base import RerankerBase
|
||||
|
||||
|
|
@ -17,6 +19,7 @@ def get_reranker() -> RerankerBase | None:
|
|||
try:
|
||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
||||
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
||||
_reranker = MxBAIReranker()
|
||||
return _reranker
|
||||
except ImportError:
|
||||
|
|
|
|||
|
|
@ -1,258 +0,0 @@
|
|||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import logfire
|
||||
from datasets import Dataset, load_dataset
|
||||
from llm_judge import ANSWER_EQUIVALENCE_RUBRIC
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
from pydantic_evals import Case
|
||||
from pydantic_evals import Dataset as EvalDataset
|
||||
from pydantic_evals.evaluators import IsInstance, LLMJudge
|
||||
from pydantic_evals.reporting import ReportCaseFailure
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
|
||||
from haiku.rag import logging # noqa
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.logging import configure_cli_logging
|
||||
from haiku.rag.qa import get_qa_agent
|
||||
|
||||
logfire.configure(send_to_logfire="if-token-present", service_name="evals")
|
||||
logfire.instrument_pydantic_ai()
|
||||
configure_cli_logging()
|
||||
console = Console()
|
||||
|
||||
QA_JUDGE_MODEL: str = "qwen3"
|
||||
db_path = Path(__file__).parent / "data" / "benchmark.lancedb"
|
||||
|
||||
|
||||
async def populate_db():
|
||||
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore
|
||||
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
|
||||
|
||||
with Progress() as progress:
|
||||
task = progress.add_task("[green]Populating database...", total=len(corpus))
|
||||
|
||||
async with HaikuRAG(db_path) as rag:
|
||||
for doc in corpus:
|
||||
uri = doc["document_id"] # type: ignore
|
||||
existing_doc = await rag.get_document_by_uri(uri)
|
||||
if existing_doc is not None:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
await rag.create_document(
|
||||
content=doc["document_extracted"], # type: ignore
|
||||
uri=uri,
|
||||
)
|
||||
progress.advance(task)
|
||||
rag.store.vacuum()
|
||||
|
||||
|
||||
async def run_match_benchmark():
|
||||
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore
|
||||
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
|
||||
|
||||
correct_at_1 = 0
|
||||
correct_at_2 = 0
|
||||
correct_at_3 = 0
|
||||
total_queries = 0
|
||||
|
||||
with Progress() as progress:
|
||||
task = progress.add_task(
|
||||
"[blue]Running retrieval benchmark...", total=len(corpus)
|
||||
)
|
||||
|
||||
async with HaikuRAG(db_path) as rag:
|
||||
for doc in corpus:
|
||||
doc_id = doc["document_id"] # type: ignore
|
||||
expected_answer = doc["answer"] # type: ignore
|
||||
if expected_answer == "The answer is not found in the document.":
|
||||
progress.advance(task)
|
||||
continue
|
||||
matches = await rag.search(
|
||||
query=doc["question"], # type: ignore
|
||||
limit=3,
|
||||
)
|
||||
|
||||
total_queries += 1
|
||||
|
||||
# Check position of correct document in results
|
||||
for position, (chunk, _) in enumerate(matches):
|
||||
assert chunk.document_id is not None, (
|
||||
"Chunk document_id should not be None"
|
||||
)
|
||||
retrieved = await rag.get_document_by_id(chunk.document_id)
|
||||
if retrieved and retrieved.uri == doc_id:
|
||||
if position == 0: # First position
|
||||
correct_at_1 += 1
|
||||
correct_at_2 += 1
|
||||
correct_at_3 += 1
|
||||
elif position == 1: # Second position
|
||||
correct_at_2 += 1
|
||||
correct_at_3 += 1
|
||||
elif position == 2: # Third position
|
||||
correct_at_3 += 1
|
||||
break
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
# Calculate recall metrics
|
||||
recall_at_1 = correct_at_1 / total_queries
|
||||
recall_at_2 = correct_at_2 / total_queries
|
||||
recall_at_3 = correct_at_3 / total_queries
|
||||
|
||||
console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Total queries: {total_queries}")
|
||||
console.print(f"Recall@1: {recall_at_1:.4f}")
|
||||
console.print(f"Recall@2: {recall_at_2:.4f}")
|
||||
console.print(f"Recall@3: {recall_at_3:.4f}")
|
||||
|
||||
return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3}
|
||||
|
||||
|
||||
async def run_qa_benchmark(k: int | None = None):
|
||||
"""Run QA benchmarking on the corpus."""
|
||||
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore
|
||||
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
|
||||
|
||||
if k is not None:
|
||||
corpus = corpus.select(range(min(k, len(corpus))))
|
||||
|
||||
cases: list[Case[str, str, dict[str, str]]] = []
|
||||
for index, doc in enumerate(corpus, start=1):
|
||||
question = doc["question"] # type: ignore[index]
|
||||
expected_answer = doc["answer"] # type: ignore[index]
|
||||
doc_id = doc["document_id"] # type: ignore[index]
|
||||
case_name = f"{index}_{doc_id}" if doc_id is not None else f"case_{index}"
|
||||
|
||||
cases.append(
|
||||
Case(
|
||||
name=case_name,
|
||||
inputs=question,
|
||||
expected_output=expected_answer,
|
||||
metadata={
|
||||
"document_id": str(doc_id),
|
||||
"case_index": str(index),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
judge_model = OpenAIChatModel(
|
||||
model_name=QA_JUDGE_MODEL,
|
||||
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
|
||||
)
|
||||
|
||||
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
|
||||
cases=cases,
|
||||
evaluators=[
|
||||
IsInstance(type_name="str"),
|
||||
LLMJudge(
|
||||
rubric=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model=judge_model,
|
||||
assertion={
|
||||
"evaluation_name": "answer_equivalent",
|
||||
"include_reason": True,
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
total_processed = 0
|
||||
passing_cases = 0
|
||||
failures: list[ReportCaseFailure[str, str, dict[str, str]]] = []
|
||||
|
||||
with Progress(console=console) as progress:
|
||||
qa_task = progress.add_task(
|
||||
"[yellow]Evaluating QA cases...",
|
||||
total=len(evaluation_dataset.cases),
|
||||
)
|
||||
|
||||
async with HaikuRAG(db_path) as rag:
|
||||
qa = get_qa_agent(rag)
|
||||
|
||||
async def answer_question(question: str) -> str:
|
||||
return await qa.answer(question)
|
||||
|
||||
for case in evaluation_dataset.cases:
|
||||
progress.console.print(f"\n[bold]Evaluating case:[/bold] {case.name}")
|
||||
|
||||
single_case_dataset = EvalDataset[str, str, dict[str, str]](
|
||||
cases=[case],
|
||||
evaluators=evaluation_dataset.evaluators,
|
||||
)
|
||||
|
||||
report = await single_case_dataset.evaluate(
|
||||
answer_question,
|
||||
name="qa_answer",
|
||||
max_concurrency=1,
|
||||
progress=False,
|
||||
)
|
||||
|
||||
total_processed += 1
|
||||
|
||||
if report.cases:
|
||||
result_case = report.cases[0]
|
||||
|
||||
equivalence = result_case.assertions.get("answer_equivalent")
|
||||
progress.console.print(f"Question: {result_case.inputs}")
|
||||
progress.console.print(f"Expected: {result_case.expected_output}")
|
||||
progress.console.print(f"Generated: {result_case.output}")
|
||||
if equivalence is not None:
|
||||
progress.console.print(
|
||||
f"Equivalent: {equivalence.value}"
|
||||
+ (f" — {equivalence.reason}" if equivalence.reason else "")
|
||||
)
|
||||
if equivalence.value:
|
||||
passing_cases += 1
|
||||
|
||||
progress.console.print("")
|
||||
|
||||
if report.failures:
|
||||
failures.extend(report.failures)
|
||||
failure = report.failures[0]
|
||||
progress.console.print(
|
||||
"[red]Failure encountered during case evaluation:[/red]"
|
||||
)
|
||||
progress.console.print(f"Question: {failure.inputs}")
|
||||
progress.console.print(f"Error: {failure.error_message}")
|
||||
progress.console.print("")
|
||||
|
||||
progress.console.print(
|
||||
f"[green]Accuracy: {(passing_cases / total_processed):.4f} "
|
||||
f"{passing_cases}/{total_processed}[/green]"
|
||||
)
|
||||
progress.advance(qa_task)
|
||||
total_cases = total_processed
|
||||
accuracy = passing_cases / total_cases if total_cases > 0 else 0
|
||||
|
||||
console.print("\n=== QA Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Total questions: {total_cases}")
|
||||
console.print(f"Correct answers: {passing_cases}")
|
||||
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
|
||||
|
||||
if failures:
|
||||
console.print("[red]\nSummary of failures:[/red]")
|
||||
for failure in failures:
|
||||
console.print(f"Case: {failure.name}")
|
||||
console.print(f"Question: {failure.inputs}")
|
||||
console.print(f"Error: {failure.error_message}")
|
||||
console.print("")
|
||||
|
||||
|
||||
async def main():
|
||||
await populate_db()
|
||||
|
||||
console.print("Running retrieval benchmarks...", style="bold blue")
|
||||
await run_match_benchmark()
|
||||
|
||||
console.print("\nRunning QA benchmarks...", style="bold yellow")
|
||||
await run_qa_benchmark()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
|
@ -1,12 +1,11 @@
|
|||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from evaluations.llm_judge import LLMJudge
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.qa.agent import QuestionAnswerAgent
|
||||
|
||||
from .llm_judge import LLMJudge
|
||||
|
||||
OPENAI_AVAILABLE = bool(Config.OPENAI_API_KEY)
|
||||
ANTHROPIC_AVAILABLE = bool(Config.ANTHROPIC_API_KEY)
|
||||
VLLM_QA_AVAILABLE = bool(Config.VLLM_QA_BASE_URL)
|
||||
|
|
|
|||
Loading…
Reference in a new issue