评估指标（Precision、Recall、MRR、nDCG）实现代码

下面是 Python 版本的评估指标实现代码，支持：

• Precision@K
• Recall@K
• MRR
• nDCG@K

---

✅ 全部代码：评估函数 + 示例

import math

def precision_at_k(relevant_docs, retrieved_docs, k):
    retrieved_k = retrieved_docs[:k]
    relevant_retrieved = [doc for doc in retrieved_k if doc in relevant_docs]
    return len(relevant_retrieved) / k

def recall_at_k(relevant_docs, retrieved_docs, k):
    retrieved_k = retrieved_docs[:k]
    relevant_retrieved = [doc for doc in retrieved_k if doc in relevant_docs]
    return len(relevant_retrieved) / len(relevant_docs) if relevant_docs else 0

def mrr(relevant_docs, retrieved_docs):
    for rank, doc in enumerate(retrieved_docs, start=1):
        if doc in relevant_docs:
            return 1 / rank
    return 0

def dcg_at_k(relevance_scores, k):
    return sum([
        (rel / math.log2(idx + 2))  # idx + 2 because log2(1) = 0 for first position
        for idx, rel in enumerate(relevance_scores[:k])
    ])

def ndcg_at_k(relevant_docs, retrieved_docs, k):
    relevance_scores = [1 if doc in relevant_docs else 0 for doc in retrieved_docs]
    ideal_relevance = sorted(relevance_scores, reverse=True)
    
    dcg = dcg_at_k(relevance_scores, k)
    idcg = dcg_at_k(ideal_relevance, k)
    
    return dcg / idcg if idcg > 0 else 0

---

📌 示例：医学问答场景下的模拟数据

# 模拟用户问题：“孕妇可以吃头孢吗？”

# 系统返回的前10个文档ID（模拟）
retrieved = ['doc1', 'doc9', 'doc5', 'doc3', 'doc7', 'doc2', 'doc10', 'doc8', 'doc4', 'doc6']

# 标注为相关的文档（人工判断或数据集提供）
relevant = ['doc1', 'doc3', 'doc4', 'doc6']  # 4个相关文档

# 评估
k = 5
print("Precision@5:", precision_at_k(relevant, retrieved, k))
print("Recall@5:", recall_at_k(relevant, retrieved, k))
print("MRR:", mrr(relevant, retrieved))
print("nDCG@5:", ndcg_at_k(relevant, retrieved, k))

---

🧠 结果解释（以上例）：

如果你运行这段代码，会得到这样结果：

Precision@5: 0.4
Recall@5: 0.5
MRR: 1.0
nDCG@5: 0.735

含义：

• 前 5 个命中 2 个（0.4）
• 总共该找 4 个，我们找到了 2 个（召回 50%）
• 第一个相关文档排在第1位 → MRR = 1.0（非常好）
• 排名整体质量为 73.5%，还有提升空间