信息检索场景下使用Python计算average precision与recall的方法

指标说明 #

• 精确率：返回结果中相关文档的占比，用于衡量检索结果的准度
• 召回率：成功检索到的相关文档占全部真实相关文档的占比，用于衡量检索结果的覆盖度
• 平均精确率（AP）：综合考虑排序质量的指标，计算每个相关文档命中位置的精确率平均值

实现代码 #

def calculate_retrieval_metrics(retrieved_docs, relevant_docs, top_k=10):
    # 转换为集合提升查找效率
    relevant_set = set(relevant_docs)
    # 截断到前top_k个返回结果
    retrieved_topk = retrieved_docs[:top_k]

    hit_num = 0
    precision_at_hit = []
    for idx, doc_id in enumerate(retrieved_topk):
        if doc_id in relevant_set:
            hit_num += 1
            # 计算当前命中位置的精确率，位置从1开始计数
            current_precision = hit_num / (idx + 1)
            precision_at_hit.append(current_precision)
    
    # 计算召回率
    recall = hit_num / len(relevant_set) if len(relevant_set) > 0 else 0.0
    # 计算平均精确率
    avg_precision = sum(precision_at_hit) / len(relevant_set) if len(relevant_set) > 0 else 0.0
    # 提取前top_k中的相关文档，按检索顺序返回
    topk_relevant = [doc for doc in retrieved_topk if doc in relevant_set]

    return {
        "average_precision": round(avg_precision, 4),
        "recall": round(recall, 4),
        "top10_relevant_docs": topk_relevant
    }

# 示例测试
if __name__ == "__main__":
    retrieved = [307, 322, 256, 325, 54, 267, 303, 333, 375, 287]
    relevant = [11, 99, 100, 307, 54]
    res = calculate_retrieval_metrics(retrieved, relevant)
    print("平均精确率：", res["average_precision"])
    print("召回率：", res["recall"])
    print("前10名中的相关文档（按检索排序）：", res["top10_relevant_docs"])

示例运行输出 #

平均精确率： 0.24
召回率： 0.4
前10名中的相关文档（按检索排序）： [307, 54]

内容的提问来源于stack exchange，提问作者Saeed Taheri