RAG召回优化完全指南：从理论到实践的三大核心策略！

from sentence_transformers import CrossEncoderimport torchfrom typing import List, Tupleimport jiebafrom sklearn.feature_extraction.text import TfidfVectorizerfrom sklearn.metrics.pairwise import cosine_similarityclass RerankingRetriever:    def __init__(self, knowledge_base: List[str], model_name: str = 'BAAI/bge-reranker-base'):        """        初始化重排序检索器                Args:            knowledge_base: 知识库文档列表            model_name: 重排序模型名称        """        self.knowledge_base = knowledge_base                # 初始化重排序模型        try:            self.reranker = CrossEncoder(model_name)            print(f"成功加载重排序模型: {model_name}")        except:            print("未能加载BGE重排序模型，将使用简化的重排序逻辑")            self.reranker = None                # 初始化TF-IDF向量器用于初步检索        self.vectorizer = TfidfVectorizer(            tokenizer=lambda x: list(jieba.cut(x)),            lowercase=False,            max_features=5000,            ngram_range=(1, 2)        )        self.doc_vectors = self.vectorizer.fit_transform(knowledge_base)        def initial_retrieval(self, query: str, top_k: int = 20) -> List[Tuple[int, float, str]]:        """        初步检索阶段                Args:            query: 查询文本            top_k: 初步检索的文档数量                    Returns:            (文档索引, 初步分数, 文档内容) 的列表        """        # TF-IDF检索        query_vector = self.vectorizer.transform([query])        similarities = cosine_similarity(query_vector, self.doc_vectors).flatten()                # 获取top_k候选        top_indices = similarities.argsort()[-top_k:][::-1]                candidates = []        for idx in top_indices:            if similarities[idx] > 0.05:  # 过滤过低相似度                candidates.append((                    idx,                     float(similarities[idx]),                     self.knowledge_base[idx]                ))                return candidates        def calculate_rerank_score(self, query: str, document: str) -> float:        """        计算重排序分数                Args:            query: 查询文本            document: 文档文本                    Returns:            重排序分数        """        if self.reranker:            # 使用BGE重排序模型            score = self.reranker.predict([(query, document)])            return float(score[0])        else:            # 简化的重排序逻辑：基于关键词匹配和长度惩罚            query_words = set(jieba.cut(query.lower()))            doc_words = set(jieba.cut(document.lower()))                        # 计算词汇重叠            overlap = len(query_words & doc_words)            union = len(query_words | doc_words)                        if union == 0:                return 0.0                        # Jaccard相似度            jaccard = overlap / union                        # 长度惩罚：过长或过短的文档得分降低            doc_len = len(document)            length_penalty = 1.0            if doc_len < 20:                length_penalty = 0.7            elif doc_len > 500:                length_penalty = 0.8                        # 查询词在文档中的位置权重            position_weight = 1.0            doc_lower = document.lower()            for word in query_words:                if word in doc_lower:                    pos = doc_lower.find(word) / len(doc_lower)                    # 越靠前权重越高                    position_weight += (1 - pos) * 0.1                        return jaccard * length_penalty * position_weight        def rerank_candidates(self, query: str, candidates: List[Tuple[int, float, str]]) -> List[dict]:        """        对候选文档进行重排序                Args:            query: 查询文本            candidates: 候选文档列表                    Returns:            重排序后的结果        """        reranked_results = []                print(f"正在对{len(candidates)}个候选文档进行重排序...")                for doc_idx, initial_score, content in candidates:            # 计算重排序分数            rerank_score = self.calculate_rerank_score(query, content)                        reranked_results.append({                'document_id': doc_idx,                'content': content,                'initial_score': initial_score,                'rerank_score': rerank_score,                'final_score': rerank_score  # 可以结合initial_score            })                # 按重排序分数排序        reranked_results.sort(key=lambda x: x['rerank_score'], reverse=True)                return reranked_results        def retrieve(self, query: str, top_k: int = 5) -> List[dict]:        """        完整的检索流程：初步检索 + 重排序                Args:            query: 查询文本            top_k: 最终返回的文档数量                    Returns:            最终检索结果        """        print(f"查询: {query}")                # 步骤1: 初步检索        print("步骤1: 初步检索...")        candidates = self.initial_retrieval(query, top_k * 4)  # 多检索一些用于重排        print(f"初步检索到{len(candidates)}个候选文档")                if not candidates:            return []                # 步骤2: 重排序        print("步骤2: 重排序...")        reranked_results = self.rerank_candidates(query, candidates)                # 返回top_k结果        final_results = reranked_results[:top_k]                print(f"最终返回{len(final_results)}个结果")        return final_results# 使用示例和对比演示def demo_reranking():    """重排序检索演示"""        # 扩展的知识库    knowledge_base = [        "深度学习模型训练效率优化的关键技术包括混合精度训练、梯度累积、模型并行等方法，可以显著提升训练速度。AdamW优化器和学习率调度也很重要。",        "深度学习的基础理论涉及神经网络的反向传播算法、激活函数的选择以及损失函数的设计原理。这些是理解深度学习的基础。",        "混合精度训练是一种重要的优化技术，通过FP16和FP32的结合使用，在保持模型精度的同时减少显存占用并加速训练过程。",        "分布式训练技术如数据并行和模型并行可以充分利用多GPU资源，大幅缩短大模型的训练时间。这对大规模模型训练至关重要。",        "PyTorch和TensorFlow是两个主流的深度学习框架，各有优势。PyTorch更适合研究和原型开发，TensorFlow更适合生产环境。",        "Transformer架构revolutionized了自然语言处理领域，其自注意力机制成为现代大语言模型如GPT和BERT的基础架构。",        "卷积神经网络CNN在计算机视觉任务中表现出色，特别适用于图像分类、目标检测和图像分割等任务。",        "优化器的选择对深度学习训练效果有重大impact。AdamW相比传统SGD在处理稀疏梯度时有明显优势，特别适合Transformer模型。",        "梯度累积技术允许在内存受限的情况下模拟大批量训练，这对训练大型模型很有帮助。通过累积多个小批量的梯度再更新参数。",        "学习率调度策略如余弦退火、warmup等对训练稳定性和最终效果有重要影响。合适的学习率调度可以显著提升模型性能。"    ]        # 初始化检索器    retriever = RerankingRetriever(knowledge_base)        # 执行检索    query = "如何提高深度学习训练效率？"    results = retriever.retrieve(query, top_k=5)        # 展示结果    print(f"\n=== 重排序检索结果 ===")    for i, result in enumerate(results, 1):        print(f"\n第{i}个结果:")        print(f"重排序分数: {result['rerank_score']:.4f}")        print(f"初步检索分数: {result['initial_score']:.4f}")        print(f"内容: {result['content']}")        print("-" * 50)if __name__ == "__main__":    demo_reranking()

import openaiimport pickleimport osfrom typing import List, Dict, Tupleimport jiebafrom sklearn.feature_extraction.text import TfidfVectorizerfrom sklearn.metrics.pairwise import cosine_similarityimport numpy as npclass QueryDocExpansionRetriever:    def __init__(self, api_key: str, knowledge_base: List[str], cache_file: str = "doc2query_cache.pkl"):        """        初始化查询-文档双向扩展检索器                Args:            api_key: OpenAI API密钥            knowledge_base: 知识库文档列表            cache_file: Doc2Query结果缓存文件        """        openai.api_key = api_key        self.knowledge_base = knowledge_base        self.cache_file = cache_file                # 加载或生成Doc2Query扩展        self.expanded_docs = self._load_or_generate_doc2query()                # 构建扩展文档的向量索引        self.vectorizer = TfidfVectorizer(            tokenizer=lambda x: list(jieba.cut(x)),            lowercase=False,            max_features=8000,            ngram_range=(1, 2)        )        self.doc_vectors = self.vectorizer.fit_transform(self.expanded_docs)                print(f"知识库初始化完成，共{len(self.knowledge_base)}个文档")        def _load_or_generate_doc2query(self) -> List[str]:        """        加载或生成Doc2Query扩展文档        """        if os.path.exists(self.cache_file):            print("从缓存加载Doc2Query扩展...")            with open(self.cache_file, 'rb') as f:                return pickle.load(f)        else:            print("生成Doc2Query扩展...")            expanded_docs = self._generate_doc2query_expansion()            # 保存到缓存            with open(self.cache_file, 'wb') as f:                pickle.dump(expanded_docs, f)            return expanded_docs        def _generate_doc2query_expansion(self) -> List[str]:        """        为每个文档生成可能的查询问题（Doc2Query）        """        expanded_docs = []                for i, doc in enumerate(self.knowledge_base):            print(f"处理文档 {i+1}/{len(self.knowledge_base)}")                        # 生成该文档可能回答的问题            generated_queries = self._generate_queries_for_doc(doc)                        # 将原文档与生成的问题拼接            expanded_content = doc            if generated_queries:                expanded_content += "\n\n可能相关的问题：\n" + "\n".join(generated_queries)                        expanded_docs.append(expanded_content)                return expanded_docs        def _generate_queries_for_doc(self, document: str, num_queries: int = 3) -> List[str]:        """        为单个文档生成可能的查询问题                Args:            document: 文档内容            num_queries: 生成问题数量                    Returns:            生成的问题列表        """        prompt = f"""        基于以下文档内容，生成{num_queries}个用户可能会问的、该文档能够回答的问题。        要求：        1. 问题要自然、具体        2. 问题应该能够通过文档内容得到答案        3. 问题表述要多样化        4. 每个问题占一行                文档内容：        {document}                生成的问题：        """                try:            response = openai.ChatCompletion.create(                model="gpt-3.5-turbo",                messages=[{"role": "user", "content": prompt}],                temperature=0.7,                max_tokens=200            )                        questions = response.choices[0].message.content.strip().split('\n')            questions = [q.strip() for q in questions if q.strip()]                        # 清理问题格式            clean_questions = []            for q in questions:                # 移除序号                q = q.strip()                q = q.lstrip('0123456789.-) ')                if len(q) > 5 and q.endswith('?'):                    clean_questions.append(q)                        return clean_questions[:num_queries]                    except Exception as e:            print(f"生成问题失败: {e}")            return []        def query2doc_expansion(self, query: str) -> str:        """        Query2Doc扩展：根据查询生成假设的文档内容                Args:            query: 原始查询                    Returns:            扩展后的查询内容        """        prompt = f"""        基于以下查询，生成一个假设的、相关的文档片段，这个文档片段应该能够回答这个查询。        要求：        1. 内容要专业、准确        2. 包含相关的关键概念和术语        3. 长度控制在100-200字        4. 不要生成过于具体的数字或引用                查询：{query}                假设的相关文档内容：        """                try:            response = openai.ChatCompletion.create(                model="gpt-3.5-turbo",                messages=[{"role": "user", "content": prompt}],                temperature=0.6,                max_tokens=250            )                        pseudo_doc = response.choices[0].message.content.strip()            return pseudo_doc                    except Exception as e:            print(f"Query2Doc扩展失败: {e}")            return ""        def retrieve_with_expansion(self, query: str, top_k: int = 5, use_query2doc: bool = True) -> List[dict]:        """        使用查询和文档扩展进行检索                Args:            query: 原始查询            top_k: 返回结果数量            use_query2doc: 是否使用Query2Doc扩展                    Returns:            检索结果列表        """        print(f"原始查询: {query}")                # Query2Doc扩展        expanded_query = query        if use_query2doc:            pseudo_doc = self.query2doc_expansion(query)            if pseudo_doc:                expanded_query = f"{query}\n\n相关内容：{pseudo_doc}"                print(f"Query2Doc扩展: {pseudo_doc}")                # 检索扩展后的文档集合        query_vector = self.vectorizer.transform([expanded_query])        similarities = cosine_similarity(query_vector, self.doc_vectors).flatten()                # 获取top_k结果        top_indices = similarities.argsort()[-top_k * 2:][::-1]  # 多取一些候选                results = []        for idx in top_indices:            if similarities[idx] > 0.1:  # 过滤低相似度结果                # 返回原始文档内容，而不是扩展后的                results.append({                    'document_id': idx,                    'content': self.knowledge_base[idx],  # 原始文档                    'expanded_content': self.expanded_docs[idx],  # 扩展文档                    'score': float(similarities[idx])                })                return results[:top_k]        def compare_methods(self, query: str, top_k: int = 3):        """        对比不同检索方法的效果                Args:            query: 查询文本            top_k: 返回结果数量        """        print(f"=== 检索方法对比 ===")        print(f"查询: {query}\n")                # 方法1：基础检索（无扩展）        print("【方法1：基础TF-IDF检索】")        basic_vectorizer = TfidfVectorizer(            tokenizer=lambda x: list(jieba.cut(x)),            lowercase=False        )        basic_vectors = basic_vectorizer.fit_transform(self.knowledge_base)        query_vec = basic_vectorizer.transform([query])        basic_similarities = cosine_similarity(query_vec, basic_vectors).flatten()        basic_top = basic_similarities.argsort()[-top_k:][::-1]                for i, idx in enumerate(basic_top, 1):            print(f"  {i}. (分数: {basic_similarities[idx]:.3f}) {self.knowledge_base[idx][:100]}...")                # 方法2：仅Doc2Query扩展        print(f"\n【方法2：Doc2Query扩展检索】")        doc2query_vec = self.vectorizer.transform([query])        doc2query_similarities = cosine_similarity(doc2query_vec, self.doc_vectors).flatten()        doc2query_top = doc2query_similarities.argsort()[-top_k:][::-1]                for i, idx in enumerate(doc2query_top, 1):            print(f"  {i}. (分数: {doc2query_similarities[idx]:.3f}) {self.knowledge_base[idx][:100]}...")                # 方法3：Query2Doc + Doc2Query双向扩展        print(f"\n【方法3：Query2Doc + Doc2Query双向扩展】")        results = self.retrieve_with_expansion(query, top_k, use_query2doc=True)                for i, result in enumerate(results, 1):            print(f"  {i}. (分数: {result['score']:.3f}) {result['content'][:100]}...")# 使用示例和完整演示def demo_expansion_retrieval():    """查询-文档双向扩展检索演示"""        # 构建更丰富的知识库    knowledge_base = [        "深度学习模型训练效率优化包括混合精度训练、梯度累积、模型并行等技术。混合精度使用FP16计算可以减少显存占用并加速训练。",        "AdamW优化器结合了Adam的自适应学习率和权重衰减正则化，在Transformer模型训练中表现优异，特别适合大模型训练。",        "分布式训练通过数据并行将批次分配到多个GPU，模型并行将大模型拆分到多个设备，可以显著缩短训练时间。",        "学习率调度策略对训练效果至关重要。常用的有余弦退火、线性衰减、warmup等，需要根据具体任务调整。",        "梯度累积技术允许在显存有限的情况下模拟大批量训练，通过累积多个小批次的梯度再进行参数更新。",        "Transformer架构使用自注意力机制处理序列数据，已成为NLP和多模态任务的主流架构，GPT和BERT都基于此架构。",        "卷积神经网络通过卷积层提取图像特征，在计算机视觉任务中应用广泛，包括图像分类、目标检测、语义分割等。",        "深度学习框架PyTorch提供动态计算图和直观的API，适合研究和原型开发，而TensorFlow更适合生产部署。",        "正则化技术如Dropout、BatchNorm、LayerNorm可以防止过拟合并提升模型泛化能力，是深度学习中的重要技术。",        "预训练模型如BERT、GPT通过大规模无监督学习获得通用语言表示，然后在下游任务上微调，大幅提升了NLP任务效果。"    ]        # 初始化检索器（需要OpenAI API Key）    api_key = "your-openai-api-key"  # 请替换为实际的API Key    retriever = QueryDocExpansionRetriever(api_key, knowledge_base)        # 测试查询    test_queries = [        "如何加速深度学习模型训练？",        "什么是混合精度训练？",        "Transformer架构有什么特点？"    ]        for query in test_queries:        print(f"\n{'='*60}")        retriever.compare_methods(query, top_k=3)        print(f"{'='*60}")if __name__ == "__main__":    demo_expansion_retrieval()

class AdvancedRAGRetriever:    def __init__(self, api_key: str, knowledge_base: List[str]):        """        高级RAG检索器，组合多种优化策略        """        self.multi_query_retriever = MultiQueryRetriever(api_key, knowledge_base)        self.reranking_retriever = RerankingRetriever(knowledge_base)        self.expansion_retriever = QueryDocExpansionRetriever(api_key, knowledge_base)        def advanced_retrieve(self, query: str, top_k: int = 5) -> List[dict]:        """        高级检索：组合多种策略                策略组合：        1. 多查询扩展增加召回多样性        2. Query2Doc丰富查询语义        3. Doc2Query增强文档匹配能力        4. 重排序精确排序        """        print(f"=== 高级RAG检索 ===")        print(f"查询: {query}")                # 步骤1: 多查询检索        print("\n步骤1: 多查询检索...")        multi_results = self.multi_query_retriever.retrieve(query, top_k * 2)                # 步骤2: 扩展检索        print("\n步骤2: 双向扩展检索...")        expansion_results = self.expansion_retriever.retrieve_with_expansion(query, top_k * 2)                # 步骤3: 合并候选结果        print("\n步骤3: 合并候选结果...")        all_candidates = {}                # 添加多查询结果        for result in multi_results:            doc_id = result['document_id']            all_candidates[doc_id] = {                'content': result['content'],                'multi_query_score': result['score'],                'expansion_score': 0.0            }                # 添加扩展检索结果        for result in expansion_results:            doc_id = result['document_id']            if doc_id in all_candidates:                all_candidates[doc_id]['expansion_score'] = result['score']            else:                all_candidates[doc_id] = {                    'content': result['content'],                    'multi_query_score': 0.0,                    'expansion_score': result['score']                }                # 步骤4: 重排序        print("\n步骤4: 重排序最终结果...")        candidates_list = [            (doc_id, max(scores['multi_query_score'], scores['expansion_score']), scores['content'])            for doc_id, scores in all_candidates.items()        ]                # 使用重排序器        reranked_results = self.reranking_retriever.rerank_candidates(query, candidates_list)                # 添加组合分数        for result in reranked_results:            doc_id = result['document_id']            if doc_id in all_candidates:                result['multi_query_score'] = all_candidates[doc_id]['multi_query_score']                result['expansion_score'] = all_candidates[doc_id]['expansion_score']                # 组合最终分数                result['combined_score'] = (                    result['rerank_score'] * 0.5 +                    result['multi_query_score'] * 0.3 +                    result['expansion_score'] * 0.2                )                # 按组合分数重新排序        reranked_results.sort(key=lambda x: x['combined_score'], reverse=True)                return reranked_results[:top_k]# 完整演示def demo_advanced_rag():    """高级RAG系统演示"""    knowledge_base = [        # ... (使用前面定义的知识库)    ]        api_key = "your-openai-api-key"    advanced_retriever = AdvancedRAGRetriever(api_key, knowledge_base)        query = "如何优化大模型训练效率？"    results = advanced_retriever.advanced_retrieve(query, top_k=3)        print(f"\n=== 最终检索结果 ===")    for i, result in enumerate(results, 1):        print(f"\n第{i}个结果:")        print(f"  组合分数: {result['combined_score']:.4f}")        print(f"  重排序分数: {result['rerank_score']:.4f}")        print(f"  多查询分数: {result['multi_query_score']:.4f}")        print(f"  扩展检索分数: {result['expansion_score']:.4f}")        print(f"  内容: {result['content']}")

def evaluate_retrieval_quality(retriever, test_queries, ground_truth):    """    评估检索质量        Args:        retriever: 检索器实例        test_queries: 测试查询列表        ground_truth: 每个查询的相关文档ID列表            Returns:        评估指标字典    """    total_precision = 0    total_recall = 0    total_ndcg = 0        for i, query in enumerate(test_queries):        results = retriever.retrieve(query, top_k=5)        retrieved_ids = [r['document_id'] for r in results]        relevant_ids = set(ground_truth[i])                # 计算Precision@5        precision = len(set(retrieved_ids) & relevant_ids) / len(retrieved_ids)        total_precision += precision                # 计算Recall@5        recall = len(set(retrieved_ids) & relevant_ids) / len(relevant_ids)        total_recall += recall        return {        'precision': total_precision / len(test_queries),        'recall': total_recall / len(test_queries)    }