""" 检索效果评估模块 提供多种评估指标和方法，用于评估RAG系统中检索结果的质量 """ import time import json import numpy as np from typing import List, Dict, Tuple, Any, Optional, Union from dataclasses import dataclass, asdict try: from langchain_core.documents import Document except ImportError: try: from langchain_core.documents import Document except ImportError: from langchain.schema import Document from sklearn.metrics import ndcg_score, precision_score, recall_score, f1_score from sentence_transformers import SentenceTransformer, util import matplotlib.pyplot as plt import seaborn as sns import pandas as pd import torch @dataclass class RetrievalResult: """检索结果数据类""" query: str retrieved_docs: List[Document] relevant_docs: List[Document] # 真实相关的文档 retrieval_time: float scores: Optional[List[float]] = None # 检索分数 @dataclass class EvaluationMetrics: """评估指标数据类""" precision_at_k: Dict[int, float] recall_at_k: Dict[int, float] f1_at_k: Dict[int, float] map_score: float # 平均精度均值 mrr: float # 平均倒数排名 ndcg_at_k: Dict[int, float] coverage: float # 覆盖率 diversity: float # 多样性 novelty: float # 新颖性 latency: float # 平均检索延迟 class RetrievalEvaluator: """检索效果评估器""" def __init__(self, embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"): """ 初始化评估器 Args: embedding_model: 用于计算语义相似度的嵌入模型 """ self.embedding_model = SentenceTransformer(embedding_model) def evaluate_retrieval(self, results: List[RetrievalResult], k_values: List[int] = [1, 3, 5, 10]) -> EvaluationMetrics: """ 评估检索结果 Args: results: 检索结果列表 k_values: 要计算的k值列表 Returns: 评估指标 """ precision_at_k = {} recall_at_k = {} f1_at_k = {} ndcg_at_k = {} total_precision = {k: 0 for k in k_values} total_recall = {k: 0 for k in k_values} total_f1 = {k: 0 for k in k_values} total_ndcg = {k: 0 for k in k_values} all_precisions = [] all_reciprocal_ranks = [] all_latencies = [] for result in results: query = result.query retrieved_docs = result.retrieved_docs relevant_docs = result.relevant_docs retrieval_time = result.retrieval_time all_latencies.append(retrieval_time) # 获取相关文档的ID或内容 relevant_ids = set() for doc in relevant_docs: # 使用文档内容作为ID，实际应用中可以使用文档ID doc_id = doc.page_content[:50] # 使用前50个字符作为ID relevant_ids.add(doc_id) # 计算每个k值的指标 for k in k_values: retrieved_k = retrieved_docs[:k] retrieved_k_ids = set() for doc in retrieved_k: doc_id = doc.page_content[:50] retrieved_k_ids.add(doc_id) # 计算交集 intersection = len(relevant_ids.intersection(retrieved_k_ids)) # 计算Precision@K precision_k = intersection / k if k > 0 else 0 total_precision[k] += precision_k # 计算Recall@K recall_k = intersection / len(relevant_ids) if len(relevant_ids) > 0 else 0 total_recall[k] += recall_k # 计算F1@K if precision_k + recall_k > 0: f1_k = 2 * (precision_k * recall_k) / (precision_k + recall_k) else: f1_k = 0 total_f1[k] += f1_k # 计算NDCG@K if result.scores: # 创建相关性分数 (1表示相关，0表示不相关) relevance_scores = [] for doc in retrieved_k: doc_id = doc.page_content[:50] relevance = 1 if doc_id in relevant_ids else 0 relevance_scores.append(relevance) # 理想排序 (所有相关文档排在前面) ideal_relevance = sorted(relevance_scores, reverse=True) # 计算NDCG if len(relevance_scores) > 1 and sum(ideal_relevance) > 0: try: ndcg_k = ndcg_score([ideal_relevance], [relevance_scores], k=k) total_ndcg[k] += ndcg_k except: # 如果计算失败，使用简化的NDCG计算 dcg = 0 idcg = 0 for i, rel in enumerate(relevance_scores): dcg += rel / np.log2(i + 2) if rel > 0 else 0 for i, rel in enumerate(ideal_relevance): idcg += rel / np.log2(i + 2) if rel > 0 else 0 ndcg_k = dcg / idcg if idcg > 0 else 0 total_ndcg[k] += ndcg_k else: total_ndcg[k] += 1.0 # 如果没有相关文档或只有一个文档，NDCG为1 # 计算平均精度 (AP) precisions = [] for i, doc in enumerate(retrieved_docs): doc_id = doc.page_content[:50] if doc_id in relevant_ids: precision_at_i = len(relevant_ids.intersection(set( d.page_content[:50] for d in retrieved_docs[:i+1] ))) / (i + 1) precisions.append(precision_at_i) ap = sum(precisions) / len(relevant_ids) if precisions else 0 all_precisions.append(ap) # 计算倒数排名 (RR) for i, doc in enumerate(retrieved_docs): doc_id = doc.page_content[:50] if doc_id in relevant_ids: rr = 1 / (i + 1) all_reciprocal_ranks.append(rr) break else: all_reciprocal_ranks.append(0) # 计算平均指标 num_results = len(results) for k in k_values: precision_at_k[k] = total_precision[k] / num_results recall_at_k[k] = total_recall[k] / num_results f1_at_k[k] = total_f1[k] / num_results ndcg_at_k[k] = total_ndcg[k] / num_results map_score = sum(all_precisions) / num_results if all_precisions else 0 mrr = sum(all_reciprocal_ranks) / num_results if all_reciprocal_ranks else 0 latency = sum(all_latencies) / num_results if all_latencies else 0 # 计算覆盖率、多样性和新颖性 coverage = self._calculate_coverage(results) diversity = self._calculate_diversity(results) novelty = self._calculate_novelty(results) return EvaluationMetrics( precision_at_k=precision_at_k, recall_at_k=recall_at_k, f1_at_k=f1_at_k, map_score=map_score, mrr=mrr, ndcg_at_k=ndcg_at_k, coverage=coverage, diversity=diversity, novelty=novelty, latency=latency ) def _calculate_coverage(self, results: List[RetrievalResult]) -> float: """计算覆盖率 - 检索到的唯一文档数与总文档数的比例""" all_retrieved = set() all_relevant = set() for result in results: for doc in result.retrieved_docs: doc_id = doc.page_content[:50] all_retrieved.add(doc_id) for doc in result.relevant_docs: doc_id = doc.page_content[:50] all_relevant.add(doc_id) coverage = len(all_retrieved) / len(all_relevant) if all_relevant else 0 return coverage def _calculate_diversity(self, results: List[RetrievalResult]) -> float: """计算多样性 - 检索结果之间的平均语义差异""" all_similarities = [] for result in results: if len(result.retrieved_docs) < 2: continue # 获取文档嵌入 doc_texts = [doc.page_content for doc in result.retrieved_docs] embeddings = self.embedding_model.encode(doc_texts, convert_to_tensor=True) # 计算文档之间的余弦相似度 cos_sim = util.pytorch_cos_sim(embeddings, embeddings) # 获取上三角矩阵（排除对角线） upper_triangle_indices = torch.triu_indices(len(cos_sim), len(cos_sim), offset=1) similarities = cos_sim[upper_triangle_indices[0], upper_triangle_indices[1]] # 多样性 = 1 - 平均相似度 diversity = 1 - similarities.mean().item() all_similarities.append(diversity) return sum(all_similarities) / len(all_similarities) if all_similarities else 0 def _calculate_novelty(self, results: List[RetrievalResult]) -> float: """计算新颖性 - 检索结果中不重复内容的比例""" total_docs = 0 unique_docs = set() for result in results: for doc in result.retrieved_docs: total_docs += 1 doc_id = doc.page_content[:50] unique_docs.add(doc_id) novelty = len(unique_docs) / total_docs if total_docs > 0 else 0 return novelty def compare_retrievers(self, retriever_results: Dict[str, List[RetrievalResult]], k_values: List[int] = [1, 3, 5, 10]) -> Dict[str, EvaluationMetrics]: """ 比较多个检索器的性能 Args: retriever_results: 检索器名称到检索结果的映射 k_values: 要计算的k值列表 Returns: 检索器名称到评估指标的映射 """ metrics = {} for name, results in retriever_results.items(): print(f"评估检索器: {name}") metrics[name] = self.evaluate_retrieval(results, k_values) return metrics def generate_report(self, metrics: Dict[str, EvaluationMetrics], save_path: Optional[str] = None) -> str: """ 生成评估报告 Args: metrics: 检索器名称到评估指标的映射 save_path: 报告保存路径 Returns: 报告文本 """ report = [] report.append("# 检索效果评估报告\n") # 创建比较表 df_data = [] for name, metric in metrics.items(): row = {"检索器": name} row.update({ f"Precision@{k}": f"{metric.precision_at_k[k]:.4f}" for k in sorted(metric.precision_at_k.keys()) }) row.update({ f"Recall@{k}": f"{metric.recall_at_k[k]:.4f}" for k in sorted(metric.recall_at_k.keys()) }) row.update({ f"F1@{k}": f"{metric.f1_at_k[k]:.4f}" for k in sorted(metric.f1_at_k.keys()) }) row.update({ f"NDCG@{k}": f"{metric.ndcg_at_k[k]:.4f}" for k in sorted(metric.ndcg_at_k.keys()) }) row.update({ "MAP": f"{metric.map_score:.4f}", "MRR": f"{metric.mrr:.4f}", "覆盖率": f"{metric.coverage:.4f}", "多样性": f"{metric.diversity:.4f}", "新颖性": f"{metric.novelty:.4f}", "延迟(ms)": f"{metric.latency*1000:.2f}" }) df_data.append(row) df = pd.DataFrame(df_data) report.append("## 指标比较表\n") report.append(df.to_string(index=False)) report.append("\n\n") # 添加指标解释 report.append("## 指标解释\n") report.append("- **Precision@K**: 前K个结果中相关文档的比例\n") report.append("- **Recall@K**: 前K个结果中相关文档占所有相关文档的比例\n") report.append("- **F1@K**: Precision和Recall的调和平均数\n") report.append("- **NDCG@K**: 归一化折扣累积增益，考虑排序位置\n") report.append("- **MAP**: 平均精度均值，所有查询的平均精度\n") report.append("- **MRR**: 平均倒数排名，第一个相关文档排名的倒数平均值\n") report.append("- **覆盖率**: 检索到的唯一文档数与总文档数的比例\n") report.append("- **多样性**: 检索结果之间的平均语义差异\n") report.append("- **新颖性**: 检索结果中不重复内容的比例\n") report.append("- **延迟**: 平均检索时间\n") # 添加最佳检索器 report.append("## 最佳检索器\n") # 找出每个指标的最佳检索器 best_metrics = {} for metric_name in ["precision_at_5", "recall_at_5", "f1_at_5", "ndcg_at_5", "map_score", "mrr"]: best_name = max(metrics.keys(), key=lambda x: getattr(metrics[x], metric_name)) best_metrics[metric_name] = best_name report.append(f"- **{metric_name}**: {best_name}\n") report_text = "".join(report) # 保存报告 if save_path: with open(save_path, "w", encoding="utf-8") as f: f.write(report_text) print(f"报告已保存到: {save_path}") return report_text def plot_metrics_comparison(self, metrics: Dict[str, EvaluationMetrics], save_path: Optional[str] = None): """ 绘制指标比较图 Args: metrics: 检索器名称到评估指标的映射 save_path: 图表保存路径 """ # 准备数据 retriever_names = list(metrics.keys()) # 创建子图 fig, axes = plt.subplots(2, 3, figsize=(18, 12)) fig.suptitle("检索器性能比较", fontsize=16) # Precision@K ax = axes[0, 0] k_values = sorted(list(metrics[retriever_names[0]].precision_at_k.keys())) for name in retriever_names: precision_values = [metrics[name].precision_at_k[k] for k in k_values] ax.plot(k_values, precision_values, marker='o', label=name) ax.set_title("Precision@K") ax.set_xlabel("K") ax.set_ylabel("Precision") ax.legend() ax.grid(True) # Recall@K ax = axes[0, 1] for name in retriever_names: recall_values = [metrics[name].recall_at_k[k] for k in k_values] ax.plot(k_values, recall_values, marker='o', label=name) ax.set_title("Recall@K") ax.set_xlabel("K") ax.set_ylabel("Recall") ax.legend() ax.grid(True) # F1@K ax = axes[0, 2] for name in retriever_names: f1_values = [metrics[name].f1_at_k[k] for k in k_values] ax.plot(k_values, f1_values, marker='o', label=name) ax.set_title("F1@K") ax.set_xlabel("K") ax.set_ylabel("F1") ax.legend() ax.grid(True) # NDCG@K ax = axes[1, 0] for name in retriever_names: ndcg_values = [metrics[name].ndcg_at_k[k] for k in k_values] ax.plot(k_values, ndcg_values, marker='o', label=name) ax.set_title("NDCG@K") ax.set_xlabel("K") ax.set_ylabel("NDCG") ax.legend() ax.grid(True) # MAP和MRR ax = axes[1, 1] map_values = [metrics[name].map_score for name in retriever_names] mrr_values = [metrics[name].mrr for name in retriever_names] x = np.arange(len(retriever_names)) width = 0.35 ax.bar(x - width/2, map_values, width, label='MAP') ax.bar(x + width/2, mrr_values, width, label='MRR') ax.set_title("MAP和MRR") ax.set_xticks(x) ax.set_xticklabels(retriever_names) ax.legend() ax.grid(True) # 其他指标 ax = axes[1, 2] other_metrics = ['coverage', 'diversity', 'novelty'] metric_values = {metric: [] for metric in other_metrics} for name in retriever_names: for metric in other_metrics: metric_values[metric].append(getattr(metrics[name], metric)) x = np.arange(len(retriever_names)) width = 0.25 for i, metric in enumerate(other_metrics): ax.bar(x + i*width, metric_values[metric], width, label=metric) ax.set_title("其他指标") ax.set_xticks(x + width) ax.set_xticklabels(retriever_names) ax.legend() ax.grid(True) plt.tight_layout() # 保存图表 if save_path: plt.savefig(save_path, dpi=300, bbox_inches='tight') print(f"图表已保存到: {save_path}") plt.show() class RetrievalTestSet: """检索测试集""" def __init__(self, queries_file: str, documents_file: str, qrels_file: str): """ 初始化测试集 Args: queries_file: 查询文件路径，每行一个查询 documents_file: 文档文件路径，每行一个文档 qrels_file: 相关性标注文件路径，格式为: query_id,doc_id,relevance """ self.queries = self._load_queries(queries_file) self.documents = self._load_documents(documents_file) self.qrels = self._load_qrels(qrels_file) def _load_queries(self, file_path: str) -> Dict[str, str]: """加载查询""" queries = {} with open(file_path, 'r', encoding='utf-8') as f: for i, line in enumerate(f): queries[str(i)] = line.strip() return queries def _load_documents(self, file_path: str) -> Dict[str, Document]: """加载文档""" documents = {} with open(file_path, 'r', encoding='utf-8') as f: for i, line in enumerate(f): doc = Document(page_content=line.strip(), metadata={"doc_id": str(i)}) documents[str(i)] = doc return documents def _load_qrels(self, file_path: str) -> Dict[str, Dict[str, int]]: """加载相关性标注""" qrels = {} with open(file_path, 'r', encoding='utf-8') as f: for line in f: parts = line.strip().split(',') if len(parts) >= 3: query_id, doc_id, relevance = parts[0], parts[1], int(parts[2]) if query_id not in qrels: qrels[query_id] = {} qrels[query_id][doc_id] = relevance return qrels def get_retrieval_results(self, retriever, top_k: int = 10) -> List[RetrievalResult]: """ 使用指定检索器获取检索结果 Args: retriever: 检索器，需要有一个retrieve(query, top_k)方法 top_k: 返回的文档数量 Returns: 检索结果列表 """ results = [] for query_id, query_text in self.queries.items(): start_time = time.time() retrieved_docs = retriever.retrieve(query_text, top_k) retrieval_time = time.time() - start_time # 获取相关文档 relevant_docs = [] if query_id in self.qrels: for doc_id, relevance in self.qrels[query_id].items(): if relevance > 0 and doc_id in self.documents: relevant_docs.append(self.documents[doc_id]) result = RetrievalResult( query=query_text, retrieved_docs=retrieved_docs, relevant_docs=relevant_docs, retrieval_time=retrieval_time ) results.append(result) return results def create_sample_test_set(): """创建示例测试集""" # 创建示例查询 queries = [ "什么是机器学习？", "深度学习和机器学习的区别是什么？", "如何评估机器学习模型的性能？", "自然语言处理有哪些应用？", "计算机视觉的基本任务是什么？" ] # 创建示例文档 documents = [ "机器学习是人工智能的一个分支，它使计算机能够在没有明确编程的情况下学习和改进。", "深度学习是机器学习的一个子集，它使用多层神经网络来模拟人脑的工作方式。", "评估机器学习模型的常用指标包括准确率、精确率、召回率和F1分数。", "自然语言处理是计算机科学和人工智能的一个分支，专注于计算机与人类语言之间的交互。", "计算机视觉是人工智能的一个领域，训练计算机解释和理解视觉世界。", "强化学习是机器学习的一个类型，它关注软件代理应该如何在环境中采取行动以最大化累积奖励。", "数据预处理是机器学习流程中的重要步骤，包括数据清洗、特征选择和特征工程。", "过拟合是机器学习中的一个常见问题，指模型在训练数据上表现良好但在新数据上表现不佳。", "卷积神经网络（CNN）是一类深度神经网络，最常用于分析视觉图像。", "循环神经网络（RNN）是一类人工神经网络，其中节点之间的连接形成有向图沿时间序列。" ] # 创建相关性标注 qrels = { "0": {"0": 2, "1": 1, "6": 1, "7": 1}, # 什么是机器学习？ "1": {"0": 1, "1": 2, "8": 1, "9": 1}, # 深度学习和机器学习的区别 "2": {"2": 2, "7": 1}, # 如何评估机器学习模型的性能 "3": {"3": 2, "9": 1}, # 自然语言处理的应用 "4": {"4": 2, "8": 1} # 计算机视觉的基本任务 } # 保存文件 with open("sample_queries.txt", "w", encoding="utf-8") as f: for query in queries: f.write(query + "\n") with open("sample_documents.txt", "w", encoding="utf-8") as f: for doc in documents: f.write(doc + "\n") with open("sample_qrels.csv", "w", encoding="utf-8") as f: for query_id, doc_relevance in qrels.items(): for doc_id, relevance in doc_relevance.items(): f.write(f"{query_id},{doc_id},{relevance}\n") print("示例测试集已创建:") print("- sample_queries.txt: 查询文件") print("- sample_documents.txt: 文档文件") print("- sample_qrels.csv: 相关性标注文件") return RetrievalTestSet("sample_queries.txt", "sample_documents.txt", "sample_qrels.csv") if __name__ == "__main__": # 创建示例测试集 test_set = create_sample_test_set() # 创建评估器 evaluator = RetrievalEvaluator() # 这里应该使用您的实际检索器 # 以下是一个模拟的检索器，用于演示 class MockRetriever: def __init__(self, name): self.name = name def retrieve(self, query, top_k=10): # 模拟检索结果 import random all_docs = list(test_set.documents.values()) # 模拟不同质量的检索器 if self.name == "good": # 好的检索器：有更高概率返回相关文档 relevant_docs = [doc for doc in all_docs if any(keyword in doc.page_content.lower() for keyword in query.lower().split()[:2])] if relevant_docs: results = relevant_docs[:min(top_k//2, len(relevant_docs))] results += random.sample(all_docs, min(top_k-len(results), len(all_docs))) else: results = random.sample(all_docs, min(top_k, len(all_docs))) elif self.name == "medium": # 中等检索器 relevant_docs = [doc for doc in all_docs if any(keyword in doc.page_content.lower() for keyword in [query.lower().split()[0]])] if relevant_docs: results = relevant_docs[:min(top_k//3, len(relevant_docs))] results += random.sample(all_docs, min(top_k-len(results), len(all_docs))) else: results = random.sample(all_docs, min(top_k, len(all_docs))) else: # 差的检索器：随机返回 results = random.sample(all_docs, min(top_k, len(all_docs))) return results # 创建不同质量的检索器 good_retriever = MockRetriever("good") medium_retriever = MockRetriever("medium") poor_retriever = MockRetriever("poor") # 获取检索结果 good_results = test_set.get_retrieval_results(good_retriever) medium_results = test_set.get_retrieval_results(medium_retriever) poor_results = test_set.get_retrieval_results(poor_retriever) # 比较检索器 retriever_results = { "好的检索器": good_results, "中等检索器": medium_results, "差的检索器": poor_results } # 评估检索器 metrics = evaluator.compare_retrievers(retriever_results) # 生成报告 report = evaluator.generate_report(metrics, "retrieval_evaluation_report.md") print(report) # 绘制比较图 evaluator.plot_metrics_comparison(metrics, "retrieval_evaluation_comparison.png")