danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Simplify the eval harness and share the embed-fill path.

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-08-14 15:03:41 +03:00
parent f1d5918d43
commit bf94efc655
No known key found for this signature in database
12 changed files with 194 additions and 193 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

267
evaluations/evaluations/benchmark.py
View file

@ -1,8 +1,8 @@
import asyncio import asyncio
import shutil import shutil
from collections.abc import Awaitable, Callable, Mapping from collections.abc import Callable, Mapping
from pathlib import Path from pathlib import Path
from typing import Any, Literal, cast from typing import Any, Literal, NamedTuple, cast
import typer import typer
from dotenv import find_dotenv, load_dotenv from dotenv import find_dotenv, load_dotenv
@ -17,6 +17,7 @@ from evaluations.config import ConversationInput, DatasetSpec
from evaluations.datasets import DATASETS from evaluations.datasets import DATASETS
from evaluations.evaluators import ( from evaluations.evaluators import (
ANSWER_EQUIVALENCE_RUBRIC, ANSWER_EQUIVALENCE_RUBRIC,
REFUSAL_ELIGIBLE_LABELS,
REFUSAL_RUBRIC, REFUSAL_RUBRIC,
ConversationEvaluator, ConversationEvaluator,
RefusalJudge, RefusalJudge,
@ -394,6 +395,8 @@ def _attach_relevant_uris(
""" """
if spec.retrieval_loader is None or spec.retrieval_mapper is None: if spec.retrieval_loader is None or spec.retrieval_mapper is None:
return return
if not any(isinstance(case.inputs, str) for case in cases):
return
corpus = spec.retrieval_loader() corpus = spec.retrieval_loader()
if limit is not None: if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus)))) corpus = corpus.select(range(min(limit, len(corpus))))
@ -424,7 +427,7 @@ def _resolve_capability_config(
return capability_model or config.qa.model return capability_model or config.qa.model
def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | None: def _live_summary(report_cases, report_failures) -> dict[str, float | int] | None:
"""Aggregate ConversationEvaluator scores across conversations. """Aggregate ConversationEvaluator scores across conversations.
Micro rates weight every turn equally (sums across conversations); macro Micro rates weight every turn equally (sums across conversations); macro
@ -447,7 +450,7 @@ def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] |
for failure in report_failures for failure in report_failures
) )
turns_total = sum(_score(case, "turns_total") for case in scored) turns_total = sum(_score(case, "turns_total") for case in scored)
turns_judged = sum(_score(case, "turns_judged") or 0 for case in scored) turns_judged = sum(_score(case, "turns_judged") for case in scored)
turns_passed = sum(_score(case, "turns_passed") for case in scored) turns_passed = sum(_score(case, "turns_passed") for case in scored)
summary: dict[str, float | int] = { summary: dict[str, float | int] = {
"conversations": len(scored), "conversations": len(scored),
@ -475,9 +478,9 @@ def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] |
_score(case, "cited_map") for case in cited _score(case, "cited_map") for case in cited
) / len(cited) ) / len(cited)
true_refusals = sum(_score(case, "true_refusals") or 0 for case in scored) true_refusals = sum(_score(case, "true_refusals") for case in scored)
false_refusals = sum(_score(case, "false_refusals") or 0 for case in scored) false_refusals = sum(_score(case, "false_refusals") for case in scored)
unanswerable = sum(_score(case, "unanswerable_turns") or 0 for case in scored) unanswerable = sum(_score(case, "unanswerable_turns") for case in scored)
refusals = true_refusals + false_refusals refusals = true_refusals + false_refusals
summary["unanswerable_turns"] = unanswerable summary["unanswerable_turns"] = unanswerable
summary["refusals"] = refusals summary["refusals"] = refusals
@ -497,7 +500,7 @@ def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None:
for case in report_cases: for case in report_cases:
refused = case.assertions.get("refused") refused = case.assertions.get("refused")
label = (case.metadata or {}).get("answerability") label = (case.metadata or {}).get("answerability")
if refused is None or label not in ("ANSWERABLE", "UNANSWERABLE"): if refused is None or label not in REFUSAL_ELIGIBLE_LABELS:
continue continue
outcomes.append((label, bool(refused.value))) outcomes.append((label, bool(refused.value)))
if not outcomes: if not outcomes:
@ -520,18 +523,29 @@ def _filter_qa_corpus(corpus, case_ids: set[str] | None):
return corpus.filter(lambda row: row.get("id") in case_ids) return corpus.filter(lambda row: row.get("id") in case_ids)
async def run_qa_benchmark( class _QARun(NamedTuple):
cases: list[Case[Any, Any, dict[str, Any]]]
db: Path
judge_config: ModelConfig
eval_name: str
experiment_metadata: dict[str, Any]
capability_factory: CapabilityFactory
capability_model: Any
def _prepare_qa_run(
spec: DatasetSpec, spec: DatasetSpec,
config: AppConfig, config: AppConfig,
limit: int | None = None, limit: int | None,
name: str | None = None, name: str | None,
db_path: Path | None = None, db_path: Path | None,
judge_model: ModelConfig | None = None, judge_model: ModelConfig | None,
target: Target = "rag-capability", target: Target,
capability_model: ModelConfig | None = None, capability_model: ModelConfig | None,
case_ids: set[str] | None = None, case_ids: set[str] | None,
document_filter: str | None = None, document_filter: str | None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None: ) -> _QARun:
"""Shared setup for the QA runners: cases, models, name and metadata."""
corpus = spec.qa_loader() corpus = spec.qa_loader()
corpus = _filter_qa_corpus(corpus, case_ids) corpus = _filter_qa_corpus(corpus, case_ids)
if limit is not None: if limit is not None:
@ -544,7 +558,74 @@ async def run_qa_benchmark(
judge_config = judge_model or DEFAULT_JUDGE_MODEL judge_config = judge_model or DEFAULT_JUDGE_MODEL
capability_config = _resolve_capability_config(target, config, capability_model) capability_config = _resolve_capability_config(target, config, capability_model)
db = spec.db_path(db_path)
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
target=target,
capability_config=capability_config,
document_filter=document_filter,
)
experiment_metadata.update(spec.experiment_metadata or {})
return _QARun(
cases=cases,
db=spec.db_path(db_path),
judge_config=judge_config,
eval_name=eval_name,
experiment_metadata=experiment_metadata,
capability_factory=_capability_factory_for_target(target),
capability_model=get_model(capability_config, config),
)
def _print_mean_task_time(report_cases, unit: str = "case") -> None:
if not report_cases:
return
mean = sum(case.task_duration for case in report_cases) / len(report_cases)
console.print(f"Avg task time per {unit}: {mean:.2f}s")
def _print_failures(failures, show_question: bool = False) -> None:
if not failures:
return
console.print("[red]\nSummary of failures:[/red]")
for failure in failures:
console.print(f"Case: {failure.name}")
if show_question:
console.print(f"Question: {failure.inputs}")
console.print(f"Error: {failure.error_message}")
console.print("")
async def run_qa_benchmark(
spec: DatasetSpec,
config: AppConfig,
limit: int | None = None,
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "rag-capability",
capability_model: ModelConfig | None = None,
case_ids: set[str] | None = None,
document_filter: str | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
run = _prepare_qa_run(
spec,
config,
limit,
name,
db_path,
judge_model,
target,
capability_model,
case_ids,
document_filter,
)
cases, judge_config = run.cases, run.judge_config
_attach_relevant_uris(cases, spec, limit) _attach_relevant_uris(cases, spec, limit)
citation_evaluator = spec.citation_evaluator citation_evaluator = spec.citation_evaluator
@ -568,43 +649,20 @@ async def run_qa_benchmark(
] ]
if citation_evaluator is not None: if citation_evaluator is not None:
evaluators.append(citation_evaluator) evaluators.append(citation_evaluator)
if spec.evaluate_refusal: # RefusalJudge scores only cases whose metadata carries an answerability
evaluators.append( # label; on unlabeled datasets it returns no score without a judge call.
RefusalJudge( evaluators.append(
rubric=REFUSAL_RUBRIC, RefusalJudge(
model=get_model(judge_config, config), rubric=REFUSAL_RUBRIC,
assertion={"evaluation_name": "refused", "include_reason": False}, model=get_model(judge_config, config),
) assertion={"evaluation_name": "refused", "include_reason": False},
) )
)
evaluation_dataset = EvalDataset[Any, str, dict[str, Any]]( evaluation_dataset = EvalDataset[Any, str, dict[str, Any]](
name=spec.key, cases=cases, evaluators=evaluators name=spec.key, cases=cases, evaluators=evaluators
) )
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
target=target,
capability_config=capability_config,
document_filter=document_filter,
)
experiment_metadata.update(spec.experiment_metadata or {})
async def _evaluate(answer_fn: Callable[[Any], Awaitable[str]]):
return await evaluation_dataset.evaluate(
answer_fn,
name=eval_name,
max_concurrency=1,
progress=True,
metadata=experiment_metadata,
)
capability_factory = _capability_factory_for_target(target)
resolved_capability_model = get_model(capability_config, config)
async def answer_question(inputs: str | ConversationInput) -> str: async def answer_question(inputs: str | ConversationInput) -> str:
if isinstance(inputs, ConversationInput): if isinstance(inputs, ConversationInput):
question = inputs.question question = inputs.question
@ -613,11 +671,11 @@ async def run_qa_benchmark(
question = inputs question = inputs
message_history = None message_history = None
result = await run_capability_question( result = await run_capability_question(
capability_factory=capability_factory, capability_factory=run.capability_factory,
db_path=db, db_path=run.db,
config=config, config=config,
question=question, question=question,
capability_model=resolved_capability_model, capability_model=run.capability_model,
document_filter=document_filter, document_filter=document_filter,
message_history=message_history, message_history=message_history,
) )
@ -633,7 +691,13 @@ async def run_qa_benchmark(
set_eval_attribute("citation_status", result.citation_status) set_eval_attribute("citation_status", result.citation_status)
return result.answer return result.answer
report = await _evaluate(answer_question) report = await evaluation_dataset.evaluate(
answer_question,
name=run.eval_name,
max_concurrency=1,
progress=True,
metadata=run.experiment_metadata,
)
total_processed = len(report.cases) total_processed = len(report.cases)
failures = report.failures failures = report.failures
@ -660,11 +724,7 @@ async def run_qa_benchmark(
console.print(f"Total questions: {total_processed}") console.print(f"Total questions: {total_processed}")
console.print(f"Correct answers: {passing_cases}") console.print(f"Correct answers: {passing_cases}")
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
if report.cases: _print_mean_task_time(report.cases)
mean_task_time = sum(case.task_duration for case in report.cases) / len(
report.cases
)
console.print(f"Avg task time per case: {mean_task_time:.2f}s")
if citation_evaluator is not None: if citation_evaluator is not None:
score_key = citation_evaluator.get_default_evaluation_name() score_key = citation_evaluator.get_default_evaluation_name()
@ -693,26 +753,16 @@ async def run_qa_benchmark(
) )
console.print(f"Mean citations per case: {mean_citations:.2f}") console.print(f"Mean citations per case: {mean_citations:.2f}")
if spec.evaluate_refusal: if (metrics := _refusal_metrics(report.cases)) is not None:
metrics = _refusal_metrics(report.cases) precision, recall, unanswerable, refusals = metrics
if metrics is not None: console.print("\n=== Refusal vs answerability labels ===", style="bold cyan")
precision, recall, unanswerable, refusals = metrics console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}")
console.print( console.print(
"\n=== Refusal vs answerability labels ===", style="bold cyan" f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} "
) "(PARTIAL excluded)"
console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}") )
console.print(
f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} "
"(PARTIAL excluded)"
)
if failures: _print_failures(failures, show_question=True)
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 return failures[0] if failures else None
@ -727,58 +777,44 @@ async def run_live_qa_benchmark(
target: Target = "rag-capability", target: Target = "rag-capability",
capability_model: ModelConfig | None = None, capability_model: ModelConfig | None = None,
case_ids: set[str] | None = None, case_ids: set[str] | None = None,
document_filter: str | None = None,
) -> None: ) -> None:
"""Replay conversations turn by turn through one capability session. """Replay conversations turn by turn through one capability session.
One case per conversation; ``limit`` counts conversations. Answers carry One case per conversation; ``limit`` counts conversations. Answers carry
forward as real message history, so prior-turn compaction is exercised. forward as real message history, so prior-turn compaction is exercised.
""" """
corpus = spec.qa_loader() run = _prepare_qa_run(
corpus = _filter_qa_corpus(corpus, case_ids) spec,
if limit is not None: config,
corpus = corpus.select(range(min(limit, len(corpus)))) limit,
name,
cases = [ db_path,
spec.qa_case_builder(index, cast(Mapping[str, Any], doc)) judge_model,
for index, doc in enumerate(corpus, start=1) target,
] capability_model,
case_ids,
judge_config = judge_model or DEFAULT_JUDGE_MODEL document_filter,
capability_config = _resolve_capability_config(target, config, capability_model) )
db = spec.db_path(db_path)
evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]]( evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]](
name=spec.key, name=spec.key,
cases=cases, cases=run.cases,
evaluators=[ evaluators=[
ConversationEvaluator( ConversationEvaluator(
rubric=ANSWER_EQUIVALENCE_RUBRIC, rubric=ANSWER_EQUIVALENCE_RUBRIC,
model=get_model(judge_config, config), model=get_model(run.judge_config, config),
) )
], ],
) )
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
target=target,
capability_config=capability_config,
)
experiment_metadata.update(spec.experiment_metadata or {})
capability_factory = _capability_factory_for_target(target)
resolved_capability_model = get_model(capability_config, config)
async def answer_conversation(questions: list[str]) -> list[str]: async def answer_conversation(questions: list[str]) -> list[str]:
results = await run_capability_conversation( results = await run_capability_conversation(
capability_factory=capability_factory, capability_factory=run.capability_factory,
db_path=db, db_path=run.db,
config=config, config=config,
questions=list(questions), questions=list(questions),
capability_model=resolved_capability_model, capability_model=run.capability_model,
compaction=spec.compaction, compaction=spec.compaction,
) )
set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results]) set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results])
@ -793,10 +829,10 @@ async def run_live_qa_benchmark(
report = await evaluation_dataset.evaluate( report = await evaluation_dataset.evaluate(
answer_conversation, answer_conversation,
name=eval_name, name=run.eval_name,
max_concurrency=1, max_concurrency=1,
progress=True, progress=True,
metadata=experiment_metadata, metadata=run.experiment_metadata,
) )
summary = _live_summary(report.cases, report.failures) summary = _live_summary(report.cases, report.failures)
@ -849,12 +885,7 @@ async def run_live_qa_benchmark(
f"{per_turn:.2f}s per turn" f"{per_turn:.2f}s per turn"
) )
if report.failures: _print_failures(report.failures)
console.print("[red]\nSummary of failures:[/red]")
for failure in report.failures:
console.print(f"Case: {failure.name}")
console.print(f"Error: {failure.error_message}")
console.print("")
async def evaluate_dataset( async def evaluate_dataset(

6
evaluations/evaluations/capability_runner.py
View file

@ -211,8 +211,6 @@ async def run_capability_conversation(
config: AppConfig, config: AppConfig,
questions: list[str], questions: list[str],
capability_model: str | Model, capability_model: str | Model,
document_filter: str | None = None,
request_limit: int | None = None,
compaction: bool = False, compaction: bool = False,
) -> list[CapabilityRunResult]: ) -> list[CapabilityRunResult]:
"""Run a conversation's user turns sequentially through one capability. """Run a conversation's user turns sequentially through one capability.
@ -229,8 +227,8 @@ async def run_capability_conversation(
db_path, db_path,
config, config,
capability_model, capability_model,
document_filter, document_filter=None,
request_limit, request_limit=None,
compaction=compaction, compaction=compaction,
) )
history: list[ModelMessage] | None = None history: list[ModelMessage] | None = None

2
evaluations/evaluations/config.py
View file

@ -60,7 +60,6 @@ DocumentLoader = Callable[[], Dataset]
DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None] DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None]
RetrievalLoader = Callable[[], Dataset] RetrievalLoader = Callable[[], Dataset]
RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None] RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None]
QAInput = str | ConversationInput
CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]] CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]]
@ -80,7 +79,6 @@ class DatasetSpec:
document_limit: int | None = None document_limit: int | None = None
retrieval_limit: int = 5 retrieval_limit: int = 5
ingest_batch_size: int | None = None ingest_batch_size: int | None = None
evaluate_refusal: bool = False
live: bool = False live: bool = False
compaction: bool = False compaction: bool = False
experiment_metadata: dict[str, Any] | None = None experiment_metadata: dict[str, Any] | None = None

21
evaluations/evaluations/datasets/mtrag.py
View file

@ -165,7 +165,7 @@ def _task_to_record(
} }
def load_clapnq_qa() -> Dataset: def _qa_records() -> list[dict[str, Any]]:
path = _download(_GEN_TASKS_FILE) path = _download(_GEN_TASKS_FILE)
qrels = _load_qrels() qrels = _load_qrels()
records = [] records = []
@ -175,7 +175,11 @@ def load_clapnq_qa() -> Dataset:
record = _task_to_record(json.loads(line), qrels) record = _task_to_record(json.loads(line), qrels)
if record is not None: if record is not None:
records.append(record) records.append(record)
return Dataset.from_list(records) return records
def load_clapnq_qa() -> Dataset:
return Dataset.from_list(_qa_records())
def build_mtrag_case( def build_mtrag_case(
@ -225,17 +229,15 @@ def _group_conversations(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
"answerability": task["answerability"], "answerability": task["answerability"],
"multi_turn_type": task["multi_turn_type"], "multi_turn_type": task["multi_turn_type"],
"question_type": list(task["question_type"]), "question_type": list(task["question_type"]),
"relevant_uris": list(task["relevant_uris"] or []),
} }
if task["relevant_uris"]:
turn["relevant_uris"] = list(task["relevant_uris"])
turns.append(turn) turns.append(turn)
conversations.append({"id": conversation_id, "turns": turns}) conversations.append({"id": conversation_id, "turns": turns})
return conversations return conversations
def load_clapnq_conversations() -> Dataset: def load_clapnq_conversations() -> Dataset:
corpus = load_clapnq_qa() return Dataset.from_list(_group_conversations(_qa_records()))
return Dataset.from_list(_group_conversations([dict(row) for row in corpus]))
def build_mtrag_live_case( def build_mtrag_live_case(
@ -243,11 +245,7 @@ def build_mtrag_live_case(
) -> Case[list[str], list[str], dict[str, Any]]: ) -> Case[list[str], list[str], dict[str, Any]]:
questions = [turn["question"] for turn in doc["turns"]] questions = [turn["question"] for turn in doc["turns"]]
metadata_turns = [ metadata_turns = [
{ {key: value for key, value in turn.items() if key != "question"}
key: value
for key, value in turn.items()
if key != "question" and value is not None
}
for turn in doc["turns"] for turn in doc["turns"]
] ]
return Case( return Case(
@ -277,7 +275,6 @@ def _mtrag_spec(key: str, variant: str) -> DatasetSpec:
citation_evaluator=CitationMAPEvaluator(), citation_evaluator=CitationMAPEvaluator(),
retrieval_limit=10, retrieval_limit=10,
ingest_batch_size=512, ingest_batch_size=512,
evaluate_refusal=True,
experiment_metadata={"mtrag_mode": "gold_prefix"}, experiment_metadata={"mtrag_mode": "gold_prefix"},
) )

7
evaluations/evaluations/evaluators/__init__.py
View file

@ -7,12 +7,17 @@ from evaluations.evaluators.judge import (
) )
from evaluations.evaluators.map import MAPEvaluator from evaluations.evaluators.map import MAPEvaluator
from evaluations.evaluators.number_match import NumberMatchEvaluator from evaluations.evaluators.number_match import NumberMatchEvaluator
from evaluations.evaluators.refusal import REFUSAL_RUBRIC, RefusalJudge from evaluations.evaluators.refusal import (
REFUSAL_ELIGIBLE_LABELS,
REFUSAL_RUBRIC,
RefusalJudge,
)
from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator
from evaluations.evaluators.transcript import TranscriptLLMJudge from evaluations.evaluators.transcript import TranscriptLLMJudge
__all__ = [ __all__ = [
"ANSWER_EQUIVALENCE_RUBRIC", "ANSWER_EQUIVALENCE_RUBRIC",
"REFUSAL_ELIGIBLE_LABELS",
"REFUSAL_RUBRIC", "REFUSAL_RUBRIC",
"CitationMAPEvaluator", "CitationMAPEvaluator",
"ConversationEvaluator", "ConversationEvaluator",

6
evaluations/evaluations/evaluators/conversation.py
View file

@ -9,9 +9,7 @@ from pydantic_evals.evaluators.llm_as_a_judge import (
) )
from evaluations.evaluators.citation import average_precision from evaluations.evaluators.citation import average_precision
from evaluations.evaluators.refusal import REFUSAL_RUBRIC from evaluations.evaluators.refusal import REFUSAL_ELIGIBLE_LABELS, REFUSAL_RUBRIC
_REFUSAL_LABELS = ("ANSWERABLE", "UNANSWERABLE")
@dataclass @dataclass
@ -81,7 +79,7 @@ class ConversationEvaluator(Evaluator):
) )
label = turn.get("answerability") label = turn.get("answerability")
if label in _REFUSAL_LABELS: if label in REFUSAL_ELIGIBLE_LABELS:
try: try:
refused = ( refused = (
await judge_output(answer, REFUSAL_RUBRIC, self.model) await judge_output(answer, REFUSAL_RUBRIC, self.model)

20
evaluations/evaluations/evaluators/map.py
View file

@ -2,6 +2,8 @@ from dataclasses import dataclass
from pydantic_evals.evaluators import Evaluator, EvaluatorContext from pydantic_evals.evaluators import Evaluator, EvaluatorContext
from evaluations.evaluators.citation import average_precision
@dataclass @dataclass
class MAPEvaluator(Evaluator): class MAPEvaluator(Evaluator):
@ -28,22 +30,6 @@ class MAPEvaluator(Evaluator):
if ctx.metadata is None: if ctx.metadata is None:
return 0.0 return 0.0
relevant_uris = set(ctx.metadata.get("relevant_uris", [])) relevant_uris = set(ctx.metadata.get("relevant_uris", []))
retrieved_uris = ctx.output
if not relevant_uris: if not relevant_uris:
return 0.0 return 0.0
return average_precision(list(ctx.output), relevant_uris)
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

5
evaluations/evaluations/evaluators/refusal.py
View file

@ -1,4 +1,5 @@
from dataclasses import dataclass from dataclasses import dataclass
from typing import Final
from pydantic_evals.evaluators import EvaluatorContext, LLMJudge from pydantic_evals.evaluators import EvaluatorContext, LLMJudge
from pydantic_evals.evaluators.evaluator import EvaluatorOutput from pydantic_evals.evaluators.evaluator import EvaluatorOutput
@ -9,7 +10,7 @@ REFUSAL_RUBRIC = (
"substantive answer." "substantive answer."
) )
_ELIGIBLE_LABELS = ("ANSWERABLE", "UNANSWERABLE") REFUSAL_ELIGIBLE_LABELS: Final = ("ANSWERABLE", "UNANSWERABLE")
@dataclass @dataclass
@ -23,6 +24,6 @@ class RefusalJudge(LLMJudge):
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput: async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
label = (ctx.metadata or {}).get("answerability") label = (ctx.metadata or {}).get("answerability")
if label not in _ELIGIBLE_LABELS: if label not in REFUSAL_ELIGIBLE_LABELS:
return {} return {}
return await super().evaluate(ctx) return await super().evaluate(ctx)

9
evaluations/tests/test_benchmark.py
View file

@ -332,7 +332,7 @@ class TestLiveSummary:
def test_none_without_scored_cases(self) -> None: def test_none_without_scored_cases(self) -> None:
from evaluations.benchmark import _live_summary from evaluations.benchmark import _live_summary
assert _live_summary([self._case({})]) is None assert _live_summary([self._case({})], []) is None
def test_micro_rate_uses_judged_turns(self) -> None: def test_micro_rate_uses_judged_turns(self) -> None:
from evaluations.benchmark import _live_summary from evaluations.benchmark import _live_summary
@ -352,7 +352,7 @@ class TestLiveSummary:
) )
] ]
summary = _live_summary(cases) summary = _live_summary(cases, [])
assert summary is not None assert summary is not None
assert summary["micro_pass_rate"] == 1.0 assert summary["micro_pass_rate"] == 1.0
@ -784,8 +784,9 @@ class TestRunQaBenchmarkCapabilityTarget:
# (the capability manages its own client via lifespan). # (the capability manages its own client via lifespan).
mock_haiku.assert_not_called() mock_haiku.assert_not_called()
# capability model defaults to qa.model when not provided # capability model defaults to qa.model when not provided
capability_call = mock_get_model.call_args_list[-1] assert any(
assert capability_call[0][0] == AppConfig().qa.model call[0][0] == AppConfig().qa.model for call in mock_get_model.call_args_list
)
assert mock_run_capability is capability_run assert mock_run_capability is capability_run
@pytest.mark.asyncio @pytest.mark.asyncio

5
evaluations/tests/test_mtrag.py
View file

@ -139,9 +139,6 @@ class TestSpecs:
} }
assert isinstance(spec.citation_evaluator, CitationMAPEvaluator) assert isinstance(spec.citation_evaluator, CitationMAPEvaluator)
def test_refusal_evaluation_enabled(self) -> None:
assert MTRAG_CLAPNQ_SPEC.evaluate_refusal is True
class TestGenerationTasks: class TestGenerationTasks:
def test_task_to_record(self) -> None: def test_task_to_record(self) -> None:
@ -255,7 +252,7 @@ class TestLiveConversations:
} }
other_case = build_mtrag_live_case(2, conversations[1]) other_case = build_mtrag_live_case(2, conversations[1])
assert other_case.metadata is not None assert other_case.metadata is not None
assert "relevant_uris" not in other_case.metadata["turns"][0] assert other_case.metadata["turns"][0]["relevant_uris"] == []
def test_live_spec(self) -> None: def test_live_spec(self) -> None:
assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC

25
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -299,21 +299,16 @@ async def _store_documents_with_chunks(
Embeds any chunks that lack embeddings, then writes the documents, chunks, Embeds any chunks that lack embeddings, then writes the documents, chunks,
and document_items tables once apiece. Restores all tables on any failure. and document_items tables once apiece. Restores all tables on any failure.
""" """
missing = [ flat = await ensure_chunks_embedded(
chunk client._config,
for _, chunks, _ in prepared [chunk for _, chunks, _ in prepared for chunk in chunks],
for chunk in chunks client.embedder,
if chunk.embedding is None )
] embedded: list[list[Chunk]] = []
if missing: position = 0
from haiku.rag.embeddings import embed_chunks for _, chunks, _ in prepared:
embedded.append(flat[position : position + len(chunks)])
embedded_flat = await embed_chunks(missing, client.embedder, client._config) position += len(chunks)
# Assign positionally: duplicate chunk texts across documents make a
# content-keyed lookup ambiguous.
for chunk, with_embedding in zip(missing, embedded_flat):
chunk.embedding = with_embedding.embedding
embedded: list[list[Chunk]] = [chunks for _, chunks, _ in prepared]
def _extract_all_items(): def _extract_all_items():
return [extract_items("", d) for _, _, d in prepared] return [extract_items("", d) for _, _, d in prepared]

14
haiku_rag_slim/haiku/rag/client/processing.py
View file

@ -352,16 +352,10 @@ async def ensure_chunks_embedded(
embedded = await embed_chunks(chunks_to_embed, embedder, config) embedded = await embed_chunks(chunks_to_embed, embedder, config)
# Build result maintaining original order # embed_chunks preserves input order; fill positionally, since duplicate
embedded_map = {(c.content, c.order): c for c in embedded} # chunk texts across documents make a content-keyed lookup ambiguous.
result = [] filled = iter(embedded)
for ch in chunks: return [ch if ch.embedding is not None else next(filled) for ch in chunks]
if ch.embedding is not None:
result.append(ch)
else:
result.append(embedded_map[(ch.content, ch.order)])
return result
def get_extension_from_content_type_or_url(url: str, content_type: str) -> str: def get_extension_from_content_type_or_url(url: str, content_type: str) -> str: