projects/repobot/v2/eval_retrieval.py

"""用 eval_qa.jsonl 的 20 道题评估 RepoBot v2 的检索：命中率和 MRR。

    python eval_retrieval.py            不加重排
    python eval_retrieval.py --rerank   加上重排
"""
import json
import sys
from pathlib import Path

from repobot import CACHE_DIR, DOCS_DIR
from retrieval import QueryRewriter, Retriever

HERE = Path(__file__).parent
questions = [json.loads(line) for line in (HERE / "eval_qa.jsonl").read_text().splitlines()]
retriever = Retriever(DOCS_DIR, CACHE_DIR, rerank="--rerank" in sys.argv)
rewriter = QueryRewriter(CACHE_DIR)

ranks = []
for qa in questions:
    results = retriever.search(rewriter.rewrite(qa["question"]), k=5)
    rank = next((i for i, (f, t) in enumerate(results, 1) if f == qa["file"] and qa["keyword"] in t), None)
    ranks.append(rank)
    if rank is None:
        print(f"没找到：{qa['question']}（应在 {qa['file']}）")

n = len(ranks)
for top in (1, 3, 5):
    print(f"前 {top} 名命中：{sum(1 for r in ranks if r and r <= top) / n:.0%}")
print(f"MRR：{sum(1 / r for r in ranks if r) / n:.3f}")