From b5892699a0958d55afea41b6d2d4b28ff9913422 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 24 Nov 2025 14:21:23 +0200 Subject: [PATCH] Use Mean Reciprocal Rank for single document evaluation as metric. Use Mean Average Precision for variable document evaluation as metric --- CHANGELOG.md | 9 + docs/benchmarks.md | 71 +++--- evaluations/evaluations/benchmark.py | 219 ++++++++++-------- evaluations/evaluations/config.py | 2 + evaluations/evaluations/datasets/repliqa.py | 2 + evaluations/evaluations/datasets/wix.py | 2 + .../evaluations/evaluators/__init__.py | 15 ++ .../{llm_judge.py => evaluators/judge.py} | 1 - evaluations/evaluations/evaluators/map.py | 49 ++++ evaluations/evaluations/evaluators/mrr.py | 37 +++ tests/test_qa.py | 2 +- 11 files changed, 279 insertions(+), 130 deletions(-) create mode 100644 evaluations/evaluations/evaluators/__init__.py rename evaluations/evaluations/{llm_judge.py => evaluators/judge.py} (97%) create mode 100644 evaluations/evaluations/evaluators/map.py create mode 100644 evaluations/evaluations/evaluators/mrr.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a4a55c3b..4cde8102 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,9 +13,18 @@ - Tracks retrieval configuration: `retrieval_limit` for number of chunks retrieved - Tracks reranking configuration: `rerank_provider` and `rerank_model` - Enables comparison of evaluation runs with different configurations in Logfire +- **Evaluations**: Refactored retrieval evaluation to use pydantic-ai experiment framework + - New `evaluators` module with `MRREvaluator` (Mean Reciprocal Rank) and `MAPEvaluator` (Mean Average Precision) + - Retrieval benchmarks now use `Dataset.evaluate()` with full Logfire experiment tracking + - Dataset specifications now declare their retrieval evaluator (MRR for RepliQA, MAP for Wix) + - Replaced Recall@K and Success@K with industry-standard MRR and MAP metrics + - Unified evaluation framework for both retrieval and QA benchmarks ### Changed +- **Evaluations**: Renamed `--qa-limit` CLI parameter to `--limit`, now applies to both retrieval and QA benchmarks +- **Evaluations**: Retrieval evaluator selection moved from runtime logic to dataset configuration + ## [0.18.0] - 2025-11-21 ### Added diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 0a9f976f..d9f96dbe 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -25,24 +25,33 @@ You can also use command-line options: - `--skip-db` - Skip updating the evaluation database - `--skip-retrieval` - Skip retrieval benchmark - `--skip-qa` - Skip QA benchmark -- `--qa-limit N` - Limit number of QA cases to evaluate -- `--name NAME` - Override the evaluation name (defaults to `{dataset}_qa_evaluation`) +- `--limit N` - Limit number of test cases for both retrieval and QA +- `--name NAME` - Override the evaluation name (defaults to `{dataset}_retrieval_evaluation` or `{dataset}_qa_evaluation`) -## Recall +## RepliQA Retrieval -In order to calculate recall, we load the `News Stories` from `repliqa_3` (1035 documents) and index them. Subsequently, we run a search over the `question` field for each row of the dataset and check whether we match the document that answers the question. Questions for which the answer cannot be found in the documents are ignored. +We use the [RepliQA](https://huggingface.co/datasets/ServiceNow/repliqa) dataset to evaluate retrieval performance. We load the `News Stories` from `repliqa_3` (1035 documents) and index them. Subsequently, we run a search over the `question` field for each row of the dataset and check whether we match the document that answers the question. Questions for which the answer cannot be found in the documents are ignored. +For RepliQA, we use **Mean Reciprocal Rank (MRR)** as the primary metric since each query has exactly one relevant document. -The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 for the top 3 results with the "bare" default settings (Ollama `qwen3`, `mxbai-embed-large` embeddings, no reranking). +**How MRR is calculated:** +- For each query, we retrieve the top-K documents and find the rank (position) of the first relevant document +- The reciprocal rank for that query is `1/rank` (e.g., if the relevant document is at position 3, the score is 1/3 ≈ 0.333) +- If no relevant document is found in the top-K results, the score is 0 +- MRR is the mean of these reciprocal ranks across all queries +- Scores range from 0 (never found) to 1 (always found at rank 1) -| Embedding Model | Document in top 1 | Document in top 3 | Reranker | -|---------------------------------------|-------------------|-------------------|------------------------| -| Ollama / `qwen3-embedding:8b` | 0.81 | 0.95 | None | -| Ollama / `qwen3-embedding:0.6b` | 0.77 | 0.97 | None | -| Ollama / `qwen3-embedding:8b` | 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 | +**Example:** If we run 3 queries and the relevant documents are found at positions 1, 2, and not found: +- Query 1: 1/1 = 1.0 +- Query 2: 1/2 = 0.5 +- Query 3: 0 (not found) +- MRR = (1.0 + 0.5 + 0) / 3 = 0.5 + +### MRR Results + +| Embedding Model | MRR | Reranker | +|---------------------------------------|-------|------------------------| +| Ollama / `qwen3-embedding:8b` | 0.91 | - | ## Question/Answer evaluation @@ -60,7 +69,7 @@ determine whether the answer is correct. The obtained accuracy is as follows: Note the significant degradation when very small models are used such as `qwen3:0.6b`. -## Wix dataset +## Wix Retrieval 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 @@ -70,26 +79,30 @@ system handles conversational, product-specific support queries. For retrieval evaluation, 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 track two complementary metrics: +relevant passage URIs. -- **Recall@K**: Fraction of relevant documents retrieved in top K results. Measures coverage. -- **Success@K**: Fraction of queries with at least one relevant document in top K. Most relevant for RAG, where finding one good document is often sufficient. +For Wix, we use **Mean Average Precision (MAP)** as the primary metric since each query has multiple relevant documents. MAP accounts for both the presence and ranking of all relevant documents. -### Recall@K Results +**How MAP is calculated:** +- For each query, we retrieve the top-K documents and identify which ones are relevant +- For each relevant document found at position k, we calculate precision@k = (number of relevant docs in top k) / k +- Average Precision (AP) for that query is the mean of these precision values, divided by the total number of relevant documents +- MAP is the mean of AP scores across all queries +- Scores range from 0 (no relevant documents found) to 1 (all relevant documents ranked at the top) -| Embedding Model | Recall@1 | Recall@3 | Recall@5 | Reranker | -|----------------------------|----------|----------|----------|------------------------| -| `qwen3-embedding:8b` | 0.31 | 0.48 | 0.54 | None | -| `qwen3-embedding:8b` | 0.36 | 0.57 | 0.68 | `mxbai-rerank-base-v2` | -| `qwen3-embedding:8b` | 0.36 | 0.58 | 0.67 | `zeroentropy` | +**Example:** If a query has 2 relevant documents (A and B), and we retrieve 5 documents [A, X, B, Y, Z]: +- A is at position 1: precision@1 = 1/1 = 1.0 (1 relevant out of top 1) +- B is at position 3: precision@3 = 2/3 ≈ 0.667 (2 relevant out of top 3) +- AP = (1.0 + 0.667) / 2 = 0.833 +- If we had another query with AP = 0.5, then MAP = (0.833 + 0.5) / 2 = 0.667 -### Success@K Results +MAP rewards systems that rank relevant documents higher, not just finding them. -| Embedding Model | Success@1 | Success@3 | Success@5 | Reranker | -|----------------------------|-----------|-----------|-----------|------------------------| -| `qwen3-embedding:8b` | 0.36 | 0.54 | 0.62 | None | -| `qwen3-embedding:8b` | 0.42 | 0.66 | 0.76 | `mxbai-rerank-base-v2` | -| `qwen3-embedding:8b` | 0.41 | 0.66 | 0.76 | `zeroentropy` | +### MAP Results + +| Embedding Model | MAP | Reranker | +|----------------------------|-------|------------------------| +| - | - | - | ## QA Accuracy diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index be938991..10fcf28e 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -8,7 +8,7 @@ import typer from dotenv import load_dotenv from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_evals import Dataset as EvalDataset +from pydantic_evals import Case, Dataset as EvalDataset from pydantic_evals.evaluators import LLMJudge from pydantic_evals.reporting import ReportCaseFailure from rich.console import Console @@ -16,7 +16,7 @@ from rich.progress import Progress from evaluations.config import DatasetSpec from evaluations.datasets import DATASETS -from evaluations.llm_judge import ANSWER_EQUIVALENCE_RUBRIC +from evaluations.evaluators import ANSWER_EQUIVALENCE_RUBRIC from evaluations.prompts import WIX_SUPPORT_PROMPT from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, find_config_file, load_yaml_config @@ -33,6 +33,30 @@ configure_cli_logging() console = Console() +def build_experiment_metadata( + dataset_key: str, + test_cases: int, + config: AppConfig, + judge_model: str, +) -> dict[str, Any]: + """Build experiment metadata for Logfire tracking.""" + return { + "dataset": dataset_key, + "test_cases": test_cases, + "embedder_provider": config.embeddings.provider, + "embedder_model": config.embeddings.model, + "embedder_dim": config.embeddings.vector_dim, + "chunk_size": config.processing.chunk_size, + "context_chunk_radius": config.processing.context_chunk_radius, + "rerank_provider": config.reranking.provider, + "rerank_model": config.reranking.model, + "qa_provider": config.qa.provider, + "qa_model": config.qa.model, + "judge_provider": "ollama", + "judge_model": judge_model, + } + + async def populate_db(spec: DatasetSpec, config: AppConfig) -> None: spec.db_path.parent.mkdir(parents=True, exist_ok=True) corpus = spec.document_loader() @@ -68,108 +92,115 @@ async def populate_db(spec: DatasetSpec, config: AppConfig) -> None: async def run_retrieval_benchmark( - spec: DatasetSpec, config: AppConfig + spec: DatasetSpec, + config: AppConfig, + limit: int | None = None, + name: str | None = None, ) -> 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() + if limit is not None: + corpus = corpus.select(range(min(limit, len(corpus)))) - recall_totals = { - 1: 0.0, - 3: 0.0, - 5: 0.0, - } - success_totals = { - 1: 0.0, - 3: 0.0, - 5: 0.0, - } - total_queries = 0 - + cases = [] with Progress() as progress: - task = progress.add_task( - "[blue]Running retrieval benchmark...", total=len(corpus) - ) - async with HaikuRAG(spec.db_path, config=config) 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 metrics for each cutoff - 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: fraction of relevant docs retrieved - recall_totals[cutoff] += matched / len(relevant) - # Success: binary - did we get at least one relevant doc? - success_totals[cutoff] += 1.0 if matched > 0 else 0.0 - + task = progress.add_task("[blue]Building retrieval cases...", total=len(corpus)) + 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 - if total_queries == 0: + case = Case( + inputs=sample.question, + metadata={"relevant_uris": sample.expected_uris}, + ) + cases.append(case) + progress.advance(task) + + if not cases: 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 + if spec.retrieval_evaluator is None: + raise ValueError(f"No retrieval evaluator configured for dataset: {spec.key}") - success_at_1 = success_totals[1] / total_queries - success_at_3 = success_totals[3] / total_queries - success_at_5 = success_totals[5] / total_queries + evaluator = spec.retrieval_evaluator + metric_name = evaluator.__class__.__name__.replace("Evaluator", "").upper() + + dataset = EvalDataset( + cases=cases, + evaluators=[evaluator], + ) + + async with HaikuRAG(spec.db_path, config=config) as rag: + + async def retrieval_target(question: str) -> list[str]: + chunks = await rag.search(query=question, limit=5) + + seen = set() + uris = [] + for chunk, _ in chunks: + if chunk.document_id is None: + continue + doc = await rag.get_document_by_id(chunk.document_id) + if doc and doc.uri and doc.uri not in seen: + uris.append(doc.uri) + seen.add(doc.uri) + + return uris + + eval_name = name if name is not None else f"{spec.key}_retrieval_evaluation" + + experiment_metadata = build_experiment_metadata( + dataset_key=spec.key, + test_cases=len(cases), + config=config, + judge_model=QA_JUDGE_MODEL, + ) + + report = await dataset.evaluate( + retrieval_target, + name=eval_name, + max_concurrency=1, + progress=True, + metadata=experiment_metadata, + ) + + total_score = 0.0 + total_cases = 0 + for case in report.cases: + if case.scores: + for score_result in case.scores.values(): + total_score += score_result.value + total_cases += 1 + + mean_score = total_score / total_cases if total_cases > 0 else 0.0 console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan") - console.print(f"Total queries: {total_queries}") - console.print("\nRecall@K (fraction of relevant docs retrieved):") - 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}") - console.print("\nSuccess@K (queries with at least one relevant doc):") - console.print(f" Success@1: {success_at_1:.4f} ({success_at_1 * 100:.1f}%)") - console.print(f" Success@3: {success_at_3:.4f} ({success_at_3 * 100:.1f}%)") - console.print(f" Success@5: {success_at_5:.4f} ({success_at_5 * 100:.1f}%)") + console.print(f"Dataset: {spec.key}") + console.print(f"Total queries: {len(cases)}") + console.print(f"{metric_name}: {mean_score:.4f}") return { - "recall@1": recall_at_1, - "recall@3": recall_at_3, - "recall@5": recall_at_5, - "success@1": success_at_1, - "success@3": success_at_3, - "success@5": success_at_5, + metric_name.lower(): mean_score, + "queries": len(cases), } async def run_qa_benchmark( spec: DatasetSpec, config: AppConfig, - qa_limit: int | None = None, + limit: int | None = None, name: str | 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)))) + if limit is not None: + corpus = corpus.select(range(min(limit, len(corpus)))) cases = [ spec.qa_case_builder(index, cast(Mapping[str, Any], doc)) @@ -207,22 +238,12 @@ async def run_qa_benchmark( eval_name = name if name is not None else f"{spec.key}_qa_evaluation" - experiment_metadata = { - "dataset": spec.key, - "test_cases": len(cases), - "embedder_provider": config.embeddings.provider, - "embedder_model": config.embeddings.model, - "embedder_dim": config.embeddings.vector_dim, - "qa_provider": config.qa.provider, - "qa_model": config.qa.model, - "judge_provider": "ollama", - "judge_model": QA_JUDGE_MODEL, - "chunk_size": config.processing.chunk_size, - "context_chunk_radius": config.processing.context_chunk_radius, - "retrieval_limit": 3, - "rerank_provider": config.reranking.provider, - "rerank_model": config.reranking.model, - } + experiment_metadata = build_experiment_metadata( + dataset_key=spec.key, + test_cases=len(cases), + config=config, + judge_model=QA_JUDGE_MODEL, + ) report = await evaluation_dataset.evaluate( answer_question, @@ -266,7 +287,7 @@ async def evaluate_dataset( skip_db: bool, skip_retrieval: bool, skip_qa: bool, - qa_limit: int | None, + limit: int | None, name: str | None, ) -> None: if not skip_db: @@ -275,11 +296,11 @@ async def evaluate_dataset( if not skip_retrieval: console.print("Running retrieval benchmarks...", style="bold blue") - await run_retrieval_benchmark(spec, config) + await run_retrieval_benchmark(spec, config, limit=limit, name=name) if not skip_qa: console.print("\nRunning QA benchmarks...", style="bold yellow") - await run_qa_benchmark(spec, config, qa_limit=qa_limit, name=name) + await run_qa_benchmark(spec, config, limit=limit, name=name) app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.") @@ -298,8 +319,8 @@ def run( 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." + limit: int | None = typer.Option( + None, "--limit", help="Limit number of test cases for both retrieval and QA." ), name: str | None = typer.Option(None, "--name", help="Override evaluation name."), ) -> None: @@ -335,7 +356,7 @@ def run( skip_db=skip_db, skip_retrieval=skip_retrieval, skip_qa=skip_qa, - qa_limit=qa_limit, + limit=limit, name=name, ) ) diff --git a/evaluations/evaluations/config.py b/evaluations/evaluations/config.py index 8e230dda..4411081c 100644 --- a/evaluations/evaluations/config.py +++ b/evaluations/evaluations/config.py @@ -5,6 +5,7 @@ from typing import Any from datasets import Dataset from pydantic_evals import Case +from pydantic_evals.evaluators import Evaluator @dataclass @@ -39,6 +40,7 @@ class DatasetSpec: qa_case_builder: CaseBuilder retrieval_loader: RetrievalLoader | None = None retrieval_mapper: RetrievalMapper | None = None + retrieval_evaluator: Evaluator | None = None document_limit: int | None = None @property diff --git a/evaluations/evaluations/datasets/repliqa.py b/evaluations/evaluations/datasets/repliqa.py index 2fc21d0f..16ef964e 100644 --- a/evaluations/evaluations/datasets/repliqa.py +++ b/evaluations/evaluations/datasets/repliqa.py @@ -5,6 +5,7 @@ from datasets import Dataset, DatasetDict, load_dataset from pydantic_evals import Case from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample +from evaluations.evaluators import MRREvaluator def load_repliqa_corpus() -> Dataset: @@ -55,4 +56,5 @@ REPLIQ_SPEC = DatasetSpec( qa_case_builder=build_repliqa_case, retrieval_loader=load_repliqa_corpus, retrieval_mapper=map_repliqa_retrieval, + retrieval_evaluator=MRREvaluator(), ) diff --git a/evaluations/evaluations/datasets/wix.py b/evaluations/evaluations/datasets/wix.py index 112ba171..6b4358c9 100644 --- a/evaluations/evaluations/datasets/wix.py +++ b/evaluations/evaluations/datasets/wix.py @@ -6,6 +6,7 @@ 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_wix_corpus() -> Dataset: @@ -78,4 +79,5 @@ WIX_SPEC = DatasetSpec( qa_case_builder=build_wix_case, retrieval_loader=load_wix_qa, retrieval_mapper=map_wix_retrieval, + retrieval_evaluator=MAPEvaluator(), ) diff --git a/evaluations/evaluations/evaluators/__init__.py b/evaluations/evaluations/evaluators/__init__.py new file mode 100644 index 00000000..cf28d8fc --- /dev/null +++ b/evaluations/evaluations/evaluators/__init__.py @@ -0,0 +1,15 @@ +from evaluations.evaluators.judge import ( + ANSWER_EQUIVALENCE_RUBRIC, + LLMJudge, + LLMJudgeResponseSchema, +) +from evaluations.evaluators.map import MAPEvaluator +from evaluations.evaluators.mrr import MRREvaluator + +__all__ = [ + "ANSWER_EQUIVALENCE_RUBRIC", + "LLMJudge", + "LLMJudgeResponseSchema", + "MAPEvaluator", + "MRREvaluator", +] diff --git a/evaluations/evaluations/llm_judge.py b/evaluations/evaluations/evaluators/judge.py similarity index 97% rename from evaluations/evaluations/llm_judge.py rename to evaluations/evaluations/evaluators/judge.py index 9b2cc8aa..d1f7beb9 100644 --- a/evaluations/evaluations/llm_judge.py +++ b/evaluations/evaluations/evaluators/judge.py @@ -5,7 +5,6 @@ from pydantic_ai.providers.ollama import OllamaProvider from haiku.rag.config import Config -# Shared rubric/prompt for answer equivalence evaluation ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent. EVALUATION CRITERIA: diff --git a/evaluations/evaluations/evaluators/map.py b/evaluations/evaluations/evaluators/map.py new file mode 100644 index 00000000..1746f6ad --- /dev/null +++ b/evaluations/evaluations/evaluators/map.py @@ -0,0 +1,49 @@ +from dataclasses import dataclass + +from pydantic_evals.evaluators import Evaluator, EvaluatorContext + + +@dataclass +class MAPEvaluator(Evaluator): + """ + Mean Average Precision evaluator for multi-document retrieval. + + AP = (sum of P@k for each relevant doc) / total relevant docs + where P@k is precision at position k. + + Appropriate for retrieval tasks where queries have multiple relevant documents. + """ + + def evaluate(self, ctx: EvaluatorContext) -> float: + """ + Calculate average precision for a single query. + + Expected context: + - ctx.metadata['relevant_uris']: set/list of relevant document URIs + - ctx.output: list of retrieved document URIs (ordered by rank) + + Returns: + float: Average precision (0.0-1.0) + """ + if ctx.metadata is None: + return 0.0 + relevant_uris = set(ctx.metadata.get("relevant_uris", [])) + retrieved_uris = ctx.output + + if not relevant_uris: + return 0.0 + + num_relevant = len(relevant_uris) + precisions = [] + num_relevant_found = 0 + + for rank, uri in enumerate(retrieved_uris, start=1): + if uri in relevant_uris: + num_relevant_found += 1 + precision_at_k = num_relevant_found / rank + precisions.append(precision_at_k) + + if not precisions: + return 0.0 + + return sum(precisions) / num_relevant diff --git a/evaluations/evaluations/evaluators/mrr.py b/evaluations/evaluations/evaluators/mrr.py new file mode 100644 index 00000000..decbc93f --- /dev/null +++ b/evaluations/evaluations/evaluators/mrr.py @@ -0,0 +1,37 @@ +from dataclasses import dataclass + +from pydantic_evals.evaluators import Evaluator, EvaluatorContext + + +@dataclass +class MRREvaluator(Evaluator): + """ + Mean Reciprocal Rank evaluator for single-document retrieval. + + MRR = 1/rank where rank is the position of the first relevant document. + Returns 0 if no relevant document is found. + + Appropriate for retrieval tasks where each query has exactly one relevant document. + """ + + def evaluate(self, ctx: EvaluatorContext) -> float: + """ + Calculate reciprocal rank for a single query. + + Expected context: + - ctx.metadata['relevant_uris']: set/list of relevant document URIs + - ctx.output: list of retrieved document URIs (ordered by rank) + + Returns: + float: 1/rank of first relevant doc, or 0.0 if not found + """ + if ctx.metadata is None: + return 0.0 + relevant_uris = set(ctx.metadata.get("relevant_uris", [])) + retrieved_uris = ctx.output + + for rank, uri in enumerate(retrieved_uris, start=1): + if uri in relevant_uris: + return 1.0 / rank + + return 0.0 diff --git a/tests/test_qa.py b/tests/test_qa.py index 597032d9..39053d23 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -2,7 +2,7 @@ import os import pytest from datasets import Dataset -from evaluations.llm_judge import LLMJudge +from evaluations.evaluators import LLMJudge from haiku.rag.client import HaikuRAG from haiku.rag.config import Config