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

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}")

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)

# 混合检索逻辑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

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]

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