114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
"""RAG Benchmark Runner：调用检索引擎对数据集逐 Case 求值并聚合指标。
|
|
|
|
只读操作，直接复用 app.retrieval.engine 的 search()，不旁路检索链路。指标按
|
|
(mode, case, repeat) 逐样本计算，再按 mode 聚合；失败样本保留在报告中但不计入汇总，
|
|
避免异常样本污染指标。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections.abc import Callable
|
|
|
|
from app.benchmarks import metrics as m
|
|
from app.benchmarks.datasets import RAGDataset
|
|
from app.contracts import (
|
|
RAGCaseResult,
|
|
RAGDatasetCase,
|
|
RAGMetrics,
|
|
RAGRunRequest,
|
|
SearchMode,
|
|
SearchRequest,
|
|
)
|
|
from app.retrieval.engine import engine
|
|
|
|
|
|
async def run_rag(
|
|
dataset: RAGDataset,
|
|
request: RAGRunRequest,
|
|
on_case: Callable[[RAGCaseResult, int, int], None] | None = None,
|
|
) -> tuple[dict[str, RAGMetrics], list[RAGCaseResult]]:
|
|
"""执行 RAG Benchmark，返回 (按 mode 聚合的指标, 全部逐样本结果)。
|
|
|
|
on_case 在每个样本求值完成后回调 (result, done, total)，供上层更新进度与事件。
|
|
"""
|
|
total = len(request.modes) * len(dataset.cases) * request.repeat
|
|
done = 0
|
|
results: list[RAGCaseResult] = []
|
|
|
|
for mode in request.modes:
|
|
for case in dataset.cases:
|
|
for repeat in range(request.repeat):
|
|
result = await _evaluate_one(case, mode, request, repeat)
|
|
results.append(result)
|
|
done += 1
|
|
if on_case is not None:
|
|
on_case(result, done, total)
|
|
|
|
metrics_by_mode = {mode.value: _aggregate(results, mode) for mode in request.modes}
|
|
return metrics_by_mode, results
|
|
|
|
|
|
async def _evaluate_one(
|
|
case: RAGDatasetCase, mode: SearchMode, request: RAGRunRequest, repeat: int
|
|
) -> RAGCaseResult:
|
|
search_request = SearchRequest(
|
|
query=case.query,
|
|
mode=mode,
|
|
limit=request.retrieval.top_k,
|
|
include_snippet=False,
|
|
)
|
|
start = time.perf_counter()
|
|
try:
|
|
response = await engine.search(search_request)
|
|
latency_ms = (time.perf_counter() - start) * 1000.0
|
|
except Exception as exc: # 单个样本失败不中断整个 Benchmark
|
|
return RAGCaseResult(
|
|
case_id=case.case_id,
|
|
mode=mode,
|
|
repeat=repeat,
|
|
latency_ms=(time.perf_counter() - start) * 1000.0,
|
|
error=str(exc),
|
|
)
|
|
|
|
retrieved_note_ids = [item.note_id for item in response.items]
|
|
retrieved_block_ids = [item.block_id for item in response.items]
|
|
expected_notes = set(case.expected_note_ids)
|
|
expected_blocks = set(case.expected_block_ids)
|
|
k = request.retrieval.top_k
|
|
|
|
return RAGCaseResult(
|
|
case_id=case.case_id,
|
|
mode=mode,
|
|
repeat=repeat,
|
|
latency_ms=latency_ms,
|
|
retrieved_note_ids=retrieved_note_ids,
|
|
retrieved_block_ids=retrieved_block_ids,
|
|
hit_at_1=m.hit_at_k(retrieved_note_ids, expected_notes, 1),
|
|
hit_at_5=m.hit_at_k(retrieved_note_ids, expected_notes, 5),
|
|
recall=m.recall_at_k(retrieved_note_ids, expected_notes, k),
|
|
reciprocal_rank=m.reciprocal_rank(retrieved_note_ids, expected_notes),
|
|
citation_hit=m.citation_hit(retrieved_block_ids, expected_blocks),
|
|
citation_applicable=bool(case.expected_block_ids),
|
|
)
|
|
|
|
|
|
def _aggregate(cases: list[RAGCaseResult], mode: SearchMode) -> RAGMetrics:
|
|
samples = [c for c in cases if c.mode == mode]
|
|
ok = [c for c in samples if c.error is None]
|
|
if not ok:
|
|
return RAGMetrics()
|
|
|
|
latencies = [c.latency_ms for c in ok]
|
|
# citation_hit_rate 只统计声明了 expected_block_ids 的样本
|
|
citation_samples = [c for c in ok if c.citation_applicable]
|
|
return RAGMetrics(
|
|
hit_at_1=m.mean([1.0 if c.hit_at_1 else 0.0 for c in ok]),
|
|
hit_at_5=m.mean([1.0 if c.hit_at_5 else 0.0 for c in ok]),
|
|
recall_at_k=m.mean([c.recall for c in ok]),
|
|
mrr=m.mean([c.reciprocal_rank for c in ok]),
|
|
citation_hit_rate=m.mean([1.0 if c.citation_hit else 0.0 for c in citation_samples]),
|
|
p50_latency_ms=m.percentile(latencies, 50.0),
|
|
p95_latency_ms=m.percentile(latencies, 95.0),
|
|
)
|