RAG 系统召回优化实战：百万文档中提升检索速度与精度的 4 大方案

朝阳区靓仔_James

于 2025-03-20 10:00:50 发布

阅读量771

点赞数 7

文章标签： python 开发语言人工智能开源智能手机

本文链接：https://blog.csdn.net/weixin_58753619/article/details/146389043

版权

Retrieval-Augmented Generation（RAG）系统是一种结合检索和生成的技术，广泛应用于问答、对话和内容生成等场景。召回环节作为 RAG 系统的核心，直接决定了系统的检索效率和质量。在本文中，我将基于一个完整的代码示例，详细介绍如何优化 RAG 系统的召回环节，解决百万级文档规模下的速度和精度问题。优化方案包括以下四个方面：

选择适合领域的预训练嵌入模型
调整混合检索的权重参数
对关键段落进行重排序（Reranking）
使用量化技术压缩向量

以下是逐步实现的思路、代码和效果分析。

1.系统背景与挑战

假设我们有一个包含 100 万篇文档的检索系统，每篇文档平均分为 10 个片段，总计 1000 万个文档片段。我们使用 SentenceTransformer 生成嵌入向量（维度通常为 768），面临的主要挑战包括：

模型加载速度慢：嵌入模型较大，加载和推理耗时。
检索速度慢：在海量文档中计算相似度开销大。
内存占用高：1000 万个 768 维向量需要大量存储空间。目标是在保证召回质量的前提下，优化检索速度和资源占用。

2.完整代码实现

以下代码展示了如何从文档分片到优化召回的完整流程。代码基于 Python，使用了 SentenceTransformer、Faiss 等库。

import numpy as npimport timefrom typing import List, Tuplefrom sentence_transformers import SentenceTransformerimport faissimport jiebafrom sklearn.feature_extraction.text import TfidfVectorizerfrom sklearn.metrics.pairwise import cosine_similarity
# 文档类class Document:    def __init__(self, id: str, content: str):        self.id = id        self.content = content
# 文档分片def chunk_documents(documents: List[Document], chunk_size: int = 100, overlap: int = 20) -> List[Document]:    chunks = []    chunk_id = 0    for doc in documents:        words = list(jieba.cut(doc.content))        for i in range(0, len(words), chunk_size - overlap):            chunk_text = "".join(words[i:i + chunk_size])            if chunk_text:                chunks.append(Document(f"{doc.id}_chunk_{chunk_id}", chunk_text))                chunk_id += 1    return chunks
# 1. 关键词检索器class KeywordRetriever:    def __init__(self, chunks: List[Document]):        self.chunks = chunks        self.vectorizer = TfidfVectorizer(tokenizer=lambda x: list(jieba.cut(x)))        self.tfidf_matrix = self.vectorizer.fit_transform([chunk.content for chunk in chunks])
    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[Document, float]]:        start_time = time.time()        query_vector = self.vectorizer.transform([query])        similarities = cosine_similarity(query_vector, self.tfidf_matrix)[0]        top_indices = np.argsort(similarities)[::-1][:top_k]        results = [(self.chunks[idx], similarities[idx]) for idx in top_indices]        print(f"关键词检索用时: {time.time() - start_time:.4f}秒")        return results
# 2. 优化的向量检索器（使用 Faiss）class OptimizedVectorRetriever:    def __init__(self, chunks: List[Document], model_name: str = "shibing624/text2vec-base-chinese"):        self.chunks = chunks        self.model = SentenceTransformer(model_name)        start_time = time.time()        self.embeddings = self.model.encode([chunk.content for chunk in chunks])        self.dimension = self.embeddings.shape[1]
        # 使用 IVF-PQ 索引        nlist = 1000  # 聚类中心数量        m = 8         # 子量化器数量        quantizer = faiss.IndexFlatIP(self.dimension)        self.index = faiss.IndexIVFPQ(quantizer, self.dimension, nlist, m, 8)        faiss.normalize_L2(self.embeddings)        self.index.train(self.embeddings)        self.index.add(self.embeddings)        self.index.nprobe = 10        print(f"Faiss 索引构建用时: {time.time() - start_time:.4f}秒")
    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[Document, float]]:        start_time = time.time()        query_embedding = self.model.encode([query])[0].reshape(1, -1)        faiss.normalize_L2(query_embedding)        scores, indices = self.index.search(query_embedding, top_k)        results = [(self.chunks[idx], scores[0][i]) for i, idx in enumerate(indices[0])]        print(f"Faiss 检索用时: {time.time() - start_time:.4f}秒")        return results
# 3. 混合检索器class HybridRetriever:    def __init__(self, chunks: List[Document], vector_weight: float = 0.7):        self.chunks = chunks        self.keyword_retriever = KeywordRetriever(chunks)        self.vector_retriever = OptimizedVectorRetriever(chunks)        self.vector_weight = vector_weight
    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[Document, float]]:        start_time = time.time()        keyword_results = self.keyword_retriever.retrieve(query, top_k=top_k*2)        vector_results = self.vector_retriever.retrieve(query, top_k=top_k*2)
        id_to_score = {}        for doc, score in keyword_results:            id_to_score[doc.id] = (1 - self.vector_weight) * score        for doc, score in vector_results:            id_to_score[doc.id] = id_to_score.get(doc.id, 0) + self.vector_weight * score
        sorted_results = sorted(id_to_score.items(), key=lambda x: x[1], reverse=True)[:top_k]        id_to_doc = {chunk.id: chunk for chunk in self.chunks}        results = [(id_to_doc[id], score) for id, score in sorted_results]        print(f"混合检索用时: {time.time() - start_time:.4f}秒")        return results

class OptimizedRAGSystem:    def __init__(self, documents, domain_model_path=None, reranker_model_path=None, use_quantization=True, vector_weight=0.7, recall_size=100):        self.documents = documents        self.vector_weight = vector_weight        self.recall_size = recall_size        if domain_model_path:            self.embed_model = SentenceTransformer(domain_model_path)        else:            self.embed_model = SentenceTransformer("shibing624/text2vec-base-chinese")        self.keyword_retriever = KeywordRetriever(documents)        if use_quantization:            self.vector_retriever = QuantizedVectorRetriever(documents)        else:            self.vector_retriever = VectorRetriever(documents)        self.hybrid_retriever = HybridRetriever(self.keyword_retriever, self.vector_retriever, vector_weight)        self.reranker = Reranker(reranker_model_path) if reranker_model_path else Reranker()        self.retriever = TwoStageRetriever(self.hybrid_retriever, self.reranker, recall_size)    def retrieve(self, query, top_k=5):        return self.retriever.retrieve(query, top_k)

# 示例演示
# 示例文档数据documents = [    Document("doc1", "自然语言处理（NLP）是人工智能和语言学的交叉学科，研究如何让计算机理解和生成人类语言。"),    Document("doc2", "机器学习是人工智能的一个子领域，它使用统计方法让计算机系统能够从数据中学习。"),    Document("doc3", "深度学习是机器学习的一种方法，它使用多层神经网络从大规模数据中学习表示。"),    Document("doc4", "词嵌入是自然语言处理中的一种技术，它将词语映射到向量空间，使得语义相似的词在向量空间中距离较近。"),    Document("doc5", "GPT（生成式预训练变换器）是一种基于Transformer架构的大型语言模型，能够生成类似人类的文本。"),    Document("doc6", "大型语言模型（LLM）是指具有大量参数和训练数据的神经网络模型，能够理解和生成人类语言。"),    Document("doc7", "检索增强生成（RAG）是一种结合了检索系统和生成模型的方法，可以提高生成内容的准确性和可靠性。"),    Document("doc8", "语义相似度是衡量两段文本在含义上相似程度的指标，常用于信息检索和问答系统。"),    Document("doc9", "向量数据库是一种专门存储和检索向量数据的数据库系统，适用于相似性搜索和AI应用。"),    Document("doc10", "知识图谱是一种结构化知识库，以图的形式表示实体之间的关系，可以增强AI系统的推理能力。"),]
# 创建优化的RAG系统实例rag_system = OptimizedRAGSystem(    documents,               # 文档数据    domain_model_path=None,  # 如果有领域特定的模型路径，则传入路径    reranker_model_path=None,  # 如果有重排序模型路径，则传入路径    use_quantization=True,   # 是否使用量化技术    vector_weight=0.7,       # 向量检索和关键词检索的权重    recall_size=100          # 第一阶段召回的文档数量)
query = "计算机如何理解人类语言"# 执行检索top_k = 5  # 获取前5个相关文档results = rag_system.retrieve(query, top_k)# 输出检索结果print(f"查询: {query}")print("检索结果:")for doc, score in results:    print(f"得分: {score:.4f}, 内容: {doc.content}")

3.优化方案详解

3.1. 选择适合领域的预训练嵌入模型

3.1.1 原理

不同的预训练嵌入模型在特定领域的表现差异很大。例如，shibing624/text2vec-base-chinese 是通用的中文模型，而领域专用模型（如医疗领域的 medical-embeddings）可能更适合特定任务。选择合适的模型可以提升语义理解能力，从而提高召回质量。

3.1.2 实现

在代码中，我们使用 SentenceTransformer 加载模型（如 shibing624/text2vec-base-chinese）。如果需要进一步优化，可以：

评估模型：通过标注数据计算 NDCG、MRR 等指标，比较多个模型的效果。
微调模型：基于领域数据微调通用模型。

3.1.2 示例代码（评估与微调）

from sentence_transformers import SentenceTransformer, lossesfrom torch.utils.data import DataLoader
# 评估模型def evaluate_model(model_name, documents, queries, relevance):    model = SentenceTransformer(model_name)    embeddings = model.encode([doc.content for doc in documents])    # 这里需要实现评估逻辑（如 NDCG），省略具体实现
# 微调模型def finetune_model(model_name, train_data, output_path):    model = SentenceTransformer(model_name)    train_loader = DataLoader(train_data, batch_size=16, shuffle=True)    train_loss = losses.CosineSimilarityLoss(model)    model.fit(train_objectives=[(train_loader, train_loss)], epochs=3, output_path=output_path)

3.1.3 效果

质量提升：在领域数据上微调后，召回率可能提升 10%-20%。
建议：若资源有限，可直接选择开源的领域模型；若有标注数据，建议微调。

3.2 调整混合检索的权重参数

3.2.1 原理

关键词检索（如 TF-IDF）速度快但缺乏语义理解，向量检索（如 SentenceTransformer + Faiss）语义准确但计算开销大。混合检索结合两者，通过权重参数（如 vector_weight）平衡速度和精度。

3.2.2 实现

在HybridRetriever 类中，我们分别调用 KeywordRetriever 和 OptimizedVectorRetriever，然后加权融合结果：

关键词得分占比：1 - vector_weight
向量得分占比：vector_weight

3.2.3 示例代码

# 混合检索逻辑id_to_score = {}
for doc, score in keyword_results:    id_to_score[doc.id] = (1 - vector_weight) * scorefor doc, score in vector_results:    id_to_score[doc.id] = id_to_score.get(doc.id, 0) + vector_weight * score

3.2.4 优化方法

网格搜索：尝试 vector_weight 从 0 到 1 的不同值，评估 NDCG。
贝叶斯优化：使用scikit-optimize 自动寻找最佳权重。

3.2.5 效果

速度提升：关键词检索减少向量计算量，检索时间可降低 50%。
精度优化：最佳权重（如 0.7）可提升召回率 5%-10%。

3.3 对关键段落进行重排序（Reranking）

3.3.1 原理

两阶段检索策略：第一阶段快速召回大量候选结果（例如 100 个），第二阶段使用更精确的模型（如交叉编码器）对候选结果重排序。这种方法在保证效率的同时提升精度。

3.3.2 实现

由于代码中未直接实现重排序，这里补充一个示例：

from sentence_transformers import CrossEncoder
class TwoStageRetriever:    def __init__(self, retriever, reranker_model="cross-encoder/ms-marco-MiniLM-L-6-v2"):        self.retriever = retriever        self.reranker = CrossEncoder(reranker_model)
    def retrieve(self, query: str, top_k: int = 5):        candidates = self.retriever.retrieve(query, top_k=100)        pairs = [(query, doc.content) for doc, _ in candidates]        scores = self.reranker.predict(pairs)        reranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)[:top_k]        return [(doc, score) for (doc, _), score in reranked]

3.3.3 效果

质量提升：重排序后，NDCG@10 可提升 15%-25%。
速度代价：重排序增加少量延迟（几十毫秒），但整体效率仍高。

3.4. 使用量化技术压缩向量

3.4.1 原理

向量量化（如 Faiss 的 IVF-PQ）通过压缩向量减少内存占用和计算量，同时保持较高检索质量。例如，将 32 位浮点向量压缩为 8 位整数，可减少约 75% 的存储空间。

3.4.2 实现

在 OptimizedVectorRetriever 中，我们使用 Faiss 的 IndexIVFPQ：

nlist：聚类中心数量，影响索引构建时间和搜索速度。
m：子量化器数量，影响压缩率。
nprobe：搜索时检查的聚类中心数量，平衡速度和精度。

3.4.3 示例代码

# IVF-PQ 索引构建nlist = 1000m = 8quantizer = faiss.IndexFlatIP(dimension)index = faiss.IndexIVFPQ(quantizer, dimension, nlist, m, 8)index.train(embeddings)index.add(embeddings)

3.4.4效果

内存减少：对于 1000 万个 768 维向量，存储从约 28GB 降至 2-3GB。
速度提升：检索时间从秒级降至毫秒级。
精度损失：轻微（<5%），可通过调整 nprobe 补偿。

4. 性能分析

以 1000 万个文档片段为例：

原始向量检索：每查询约 2-3 秒，内存 28GB。
优化后：
- 检索时间：几十毫秒（QPS > 100）。
- 内存占用：2-3GB（压缩率 10x）。
- 召回率：提升 10%-20%（混合检索 + 重排序）。

5.实际应用建议

数据规模小：直接使用 SentenceTransformer + Faiss。
数据规模大：采用混合检索 + 量化。
实时性要求高：增加关键词检索比例，优化 nprobe。
质量要求高：引入重排序，微调嵌入模型。通过调整参数（如nlist

、vector_weight），可在速度和精度间找到最佳平衡。

6. 总结

本文从原理到代码，展示了如何优化 RAG 系统的召回环节。无论是选择领域模型、混合检索、重排序，还是向量量化，每种方法都针对特定问题提供了解决方案。在实际应用中，可根据数据规模、硬件资源和业务需求灵活组合这些技术，构建高效且准确的检索系统。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述