第 4 周：检索变好了吗？recall@k、MRR 与 nDCG

{"id": "q18", "text": "我收藏的文章列表显示错了",
 "rel": {"BUG-118": 2, "REQ-09": 2, "EXP-05": 1}}

"""检索评测：recall@k、precision@k、MRR、nDCG@k，以及两个检索器的按查询配对 bootstrap。

运行：python3 retrieval_eval.py（依赖 numpy；有 scikit-learn 时顺带用 ndcg_score 校验 nDCG）
数据：同目录 data.json，BugHunt-Bench 场景的小型自制标注集（38 篇文档、30 条查询、0/1/2 分级相关性）。
"""
import json
import math
import re
from pathlib import Path

import numpy as np

K = 5            # 汇报 @5：Agent 一次只往上下文里塞前 5 篇
DEPTH = 10       # MRR 只看前 10 名（MRR@10），10 名以外没找到记 0
N_BOOT = 10_000


# ---------- 指标：ranking 是检索结果的文档 id 列表，rel 是 {文档 id: 相关性等级} ----------

def recall_at_k(ranking, rel, k):
    relevant = {d for d, g in rel.items() if g > 0}
    return len(relevant & set(ranking[:k])) / len(relevant)


def precision_at_k(ranking, rel, k):
    return sum(rel.get(d, 0) > 0 for d in ranking[:k]) / k   # 分母固定是 k，结果不足 k 篇也一样


def reciprocal_rank(ranking, rel, depth=DEPTH):
    for i, d in enumerate(ranking[:depth], start=1):
        if rel.get(d, 0) > 0:
            return 1 / i
    return 0.0


def dcg(gains):
    return sum(g / math.log2(i + 1) for i, g in enumerate(gains, start=1))


def ndcg_at_k(ranking, rel, k):
    ideal = dcg(sorted(rel.values(), reverse=True)[:k])     # 理想排序来自全部标注，不只是检索到的
    return dcg([rel.get(d, 0) for d in ranking[:k]]) / ideal


# ---------- 两个检索器：TF-IDF 余弦 与 BM25，中文按字的一元 + 二元切分 ----------

def tokenize(text):
    tokens = []
    for run in re.findall(r"[a-z0-9]+|[一-鿿]+", text.lower()):
        if run.isascii():
            tokens.append(run)
        else:
            tokens += list(run) + [run[i:i + 2] for i in range(len(run) - 1)]
    return tokens


class Index:
    def __init__(self, docs):
        self.ids = list(docs)
        self.toks = [tokenize(docs[d]) for d in self.ids]
        self.n = len(self.ids)
        self.df = {}
        for toks in self.toks:
            for t in set(toks):
                self.df[t] = self.df.get(t, 0) + 1
        self.avgdl = sum(map(len, self.toks)) / self.n

    def _rank(self, scores):
        order = sorted(range(self.n), key=lambda i: (-scores[i], self.ids[i]))   # 同分按 id 排，结果可复现
        return [self.ids[i] for i in order]

    def tfidf(self, query):
        idf = lambda t: math.log((1 + self.n) / (1 + self.df.get(t, 0))) + 1
        def vec(toks):
            v = {}
            for t in toks:
                v[t] = v.get(t, 0) + idf(t)
            norm = math.sqrt(sum(x * x for x in v.values())) or 1.0
            return {t: x / norm for t, x in v.items()}
        q = vec(tokenize(query))
        return self._rank([sum(q.get(t, 0) * x for t, x in vec(toks).items()) for toks in self.toks])

    def bm25(self, query, k1=1.2, b=0.75):
        qt = tokenize(query)
        scores = []
        for toks in self.toks:
            tf, s = {}, 0.0
            for t in toks:
                tf[t] = tf.get(t, 0) + 1
            for t in qt:
                if t in tf:
                    idf = math.log(1 + (self.n - self.df[t] + 0.5) / (self.df[t] + 0.5))
                    s += idf * tf[t] * (k1 + 1) / (tf[t] + k1 * (1 - b + b * len(toks) / self.avgdl))
            scores.append(s)
        return self._rank(scores)


# ---------- 按查询配对：bootstrap 区间 + 随机化检验 ----------

def paired_bootstrap(a, b, n_boot=N_BOOT, seed=0):
    """a、b 是两个系统在同一批查询上的逐查询得分。整条查询一起重抽，返回 (b - a) 均值差的 95% 区间。"""
    diff = np.asarray(b) - np.asarray(a)
    idx = np.random.default_rng(seed).integers(0, diff.size, size=(n_boot, diff.size))
    return np.percentile(diff[idx].mean(axis=1), [2.5, 97.5])


def randomization_test(a, b, n_perm=N_BOOT, seed=0):
    """原假设：两个系统可以互换。每条查询随机交换 A、B（差值随机变号），看均值差至少这么大的比例（双侧）。"""
    diff = np.asarray(b) - np.asarray(a)
    signs = np.random.default_rng(seed).choice([-1, 1], size=(n_perm, diff.size))
    null = np.abs((signs * diff).mean(axis=1))
    return (np.sum(null >= abs(diff.mean()) - 1e-12) + 1) / (n_perm + 1)


def queries_needed(sd, delta, z_alpha=1.959964, z_beta=0.841621):
    """配对设计、双侧 α = 0.05、功效 80% 的正态近似：要检出均值差 delta 需要多少条查询。"""
    return math.ceil(((z_alpha + z_beta) * sd / delta) ** 2)


def self_check():
    rel = {"d1": 2, "d2": 1, "d3": 2}
    ranking = ["x", "d1", "y", "d2", "z"]
    assert recall_at_k(ranking, rel, 5) == 2 / 3
    assert precision_at_k(ranking, rel, 5) == 2 / 5
    assert reciprocal_rank(ranking, rel) == 1 / 2
    expected = (2 / math.log2(3) + 1 / math.log2(5)) / (2 + 2 / math.log2(3) + 1 / math.log2(4))
    assert abs(ndcg_at_k(ranking, rel, 5) - expected) < 1e-12
    try:  # 用 scikit-learn 交叉校验：分数给全部文档，理想排序也来自全部文档
        from sklearn.metrics import ndcg_score
        docs = ranking + ["d3"]
        y_true = [[rel.get(d, 0) for d in docs]]
        y_score = [[len(docs) - i for i in range(len(docs))]]
        assert abs(ndcg_score(y_true, y_score, k=5) - expected) < 1e-12
    except ImportError:
        pass


def main():
    self_check()
    data = json.loads(Path(__file__).with_name("data.json").read_text(encoding="utf-8"))
    index, queries = Index(data["docs"]), data["queries"]
    metrics = {
        f"recall@{K}": lambda r, rel: recall_at_k(r, rel, K),
        f"precision@{K}": lambda r, rel: precision_at_k(r, rel, K),
        f"MRR@{DEPTH}": reciprocal_rank,
        f"nDCG@{K}": lambda r, rel: ndcg_at_k(r, rel, K),
    }
    runs = {"A: TF-IDF": [index.tfidf(q["text"]) for q in queries],
            "B: BM25": [index.bm25(q["text"]) for q in queries]}
    scores = {(s, m): np.array([f(r, q["rel"]) for r, q in zip(rs, queries)])
              for s, rs in runs.items() for m, f in metrics.items()}

    print(f"{len(data['docs'])} 篇文档，{len(queries)} 条查询，bootstrap / 随机化各 {N_BOOT} 次\n")
    print(f"{'指标':<14}{'A: TF-IDF':>10}{'B: BM25':>10}{'B - A':>9}   {'配对 bootstrap 95% 区间':<24}{'随机化 p':>8}")
    for m in metrics:
        a, b = scores[("A: TF-IDF", m)], scores[("B: BM25", m)]
        lo, hi = paired_bootstrap(a, b)
        print(f"{m:<14}{a.mean():>10.3f}{b.mean():>10.3f}{b.mean() - a.mean():>+9.3f}   "
              f"[{lo:+.3f}, {hi:+.3f}]{'':<10}{randomization_test(a, b):>8.3f}")

    m = f"nDCG@{K}"
    diff = scores[("B: BM25", m)] - scores[("A: TF-IDF", m)]
    print(f"\n{m} 逐查询差值（B - A）：B 更好 {int((diff > 1e-9).sum())} 条，"
          f"A 更好 {int((diff < -1e-9).sum())} 条，打平 {int((np.abs(diff) <= 1e-9).sum())} 条")
    for q, d in sorted(zip(queries, diff), key=lambda x: x[1]):
        if abs(d) > 1e-9:
            print(f"  {q['id']} {d:+.3f}  {q['text']}")
    sd = diff.std(ddof=1)
    print(f"差值标准差 {sd:.3f}；要以 80% 功效检出 nDCG 差 0.02 / 0.05，约需 "
          f"{queries_needed(sd, 0.02)} / {queries_needed(sd, 0.05)} 条查询")


if __name__ == "__main__":
    main()

38 篇文档，30 条查询，bootstrap / 随机化各 10000 次

指标             A: TF-IDF   B: BM25    B - A   配对 bootstrap 95% 区间        随机化 p
recall@5           0.844     0.828   -0.017   [-0.042, +0.000]             0.507
precision@5        0.433     0.420   -0.013   [-0.033, +0.000]             0.507
MRR@10             0.950     1.000   +0.050   [+0.000, +0.117]             0.246
nDCG@5             0.871     0.890   +0.018   [-0.012, +0.055]             0.305

nDCG@5 逐查询差值（B - A）：B 更好 6 条，A 更好 8 条，打平 16 条
  q24 -0.109  编辑器页面提交过哪些 Bug
  q28 -0.092  接口返回 422 但是前端不显示错误
  q12 -0.080  没有头像的用户显示什么
  q20 -0.070  连点两下发布出现两篇一样的文章
  q25 -0.050  文章列表每页显示多少篇
  q01 -0.018  点了收藏数字没有加一
  q13 -0.012  筛选标签以后翻页是空白页
  q21 -0.012  改了密码以后老密码还能用
  q23 +0.022  设置页面已经测过哪些功能
  q03 +0.037  密码输错了界面上没有任何报错
  q27 +0.140  Your Feed 应该显示谁的文章
  q14 +0.236  修改文章正文后 tag 没了
  q30 +0.236  登录页和注册页测过了吗
  q18 +0.326  我收藏的文章列表显示错了
差值标准差 0.095；要以 80% 功效检出 nDCG 差 0.02 / 0.05，约需 179 / 29 条查询