构建AI智能体：RAG超越语义搜索：如何用Rerank模型实现检索精度的大幅提升

<code class="javascript">from langchain_community.vectorstores import FAISSfrom langchain.embeddings import HuggingFaceEmbeddingsfrom langchain.schema import Documentfrom transformers import AutoModelForSequenceClassification, AutoTokenizerimport torchimport numpy as np # 1. 准备示例知识库和初始检索# 初始化一个嵌入模型用于初始检索embedding_model = HuggingFaceEmbeddings(    model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2") # 假设的文档库documents = [    "苹果公司于1976年由史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗纳德·韦恩创立，最初主要生产和销售个人电脑。",    "苹果公司最著名的产品是iPhone智能手机，它彻底改变了移动通信行业。",    "水果苹果是一种蔷薇科苹果属的落叶乔木果实，营养价值高，富含维生素和纤维。",    "苹果公司在2023年发布了其首款混合现实头显设备Apple Vision Pro。",    "吃苹果有助于促进消化和增强免疫力，是一种健康零食。",    "蒂姆·库克是苹果公司的现任首席执行官，于2011年接替史蒂夫·乔布斯。"]# 转换为LangChain Document对象docs = [Document(page_content=text) for text in documents] # 构建FAISS向量库vectorstore = FAISS.from_documents(docs, embedding_model) # 用户查询query = "苹果公司的创始人是谁？" # 初始检索：使用向量库召回Top K个文档（这里K=4）initial_retrieved_docs = vectorstore.similarity_search(query, k=4)print("=== 初始检索结果（基于语义相似度）===")for i, doc in enumerate(initial_retrieved_docs):    print(f"{i+1}. [相似度得分: N/A] {doc.page_content}") # 2. 初始化Rerank模型# 我们使用 `BAAI/bge-reranker-base`，这是一个强大的中英双语Rerankermodel_name = "BAAI/bge-reranker-base"tokenizer = AutoTokenizer.from_pretrained(model_name)model = AutoModelForSequenceClassification.from_pretrained(model_name)model.eval() # 设置为评估模式 # 3. 重排序函数def rerank_docs(query, retrieved_docs, model, tokenizer, top_n=3):    """    对检索到的文档进行重排序    Args:        query: 用户查询        retrieved_docs: 检索到的文档列表        model: 重排序模型        tokenizer: 重排序模型的tokenizer        top_n: 返回顶部多少个文档    Returns:        sorted_docs: 按相关性分数降序排列的文档列表        scores: 对应的分数列表    """    # 构建（query, doc）对    pairs = [(query, doc.page_content) for doc in retrieved_docs]        #  Tokenize所有文本对    with torch.no_grad(): # 禁用梯度计算，加快推理速度        inputs = tokenizer(            pairs,             padding=True,             truncation=True,             return_tensors='pt',             max_length=512        )                # 模型前向传播，计算分数        scores = model(**inputs).logits.squeeze(dim=-1).float().numpy()     # 将分数和文档打包在一起    doc_score_list = list(zip(retrieved_docs, scores))        # 按分数降序排序    doc_score_list.sort(key=lambda x: x[1], reverse=True)        # 解包，返回Top N个文档和它们的分数    sorted_docs = [doc for doc, score in doc_score_list[:top_n]]    sorted_scores = [score for doc, score in doc_score_list[:top_n]]        return sorted_docs, sorted_scores # 4. 执行重排序reranked_docs, reranked_scores = rerank_docs(query, initial_retrieved_docs, model, tokenizer, top_n=3) # 5. 打印结果print("=== 重排序后结果（基于交叉编码相关性）===")for i, (doc, score) in enumerate(zip(reranked_docs, reranked_scores)):    print(f"{i+1}. [相关性分数: {score:.4f}] {doc.page_content}") # 6. 将精排后的上下文传递给LLM（这里用打印模拟）context_for_llm = "".join([doc.page_content for doc in reranked_docs])prompt = f"""基于以下上下文，请回答问题。上下文：{context_for_llm}=== 最终输出最优的匹配结果 ===问题：{query}答案：{reranked_docs[0].page_content}"""print(f"=== 最终提供给LLM的Prompt ===")print(prompt)</code>

<code class="javascript">=== 初始检索结果（基于语义相似度）===1. [相似度得分: N/A] 苹果公司于1976年由史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗纳德·韦恩创立，最初主要生产和销售个人电脑 。2. [相似度得分: N/A] 蒂姆·库克是苹果公司的现任首席执行官，于2011年接替史蒂夫·乔布斯。3. [相似度得分: N/A] 苹果公司在2023年发布了其首款混合现实头显设备Apple Vision Pro。4. [相似度得分: N/A] 水果苹果是一种蔷薇科苹果属的落叶乔木果实，营养价值高，富含维生素和纤维。 === 重排序后结果（基于交叉编码相关性）===1. [相关性分数: 9.3700] 苹果公司于1976年由史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗纳德·韦恩创立，最初主要生产和销售个人电脑。2. [相关性分数: 4.5305] 蒂姆·库克是苹果公司的现任首席执行官，于2011年接替史蒂夫·乔布斯。3. [相关性分数: -6.0679] 苹果公司在2023年发布了其首款混合现实头显设备Apple Vision Pro。 === 最终提供给LLM的Prompt === 基于以下上下文，请回答问题。上下文：苹果公司于1976年由史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗纳德·韦恩创立，最初主要生产和销售个人电脑。蒂姆·库克是苹果公司的现任首席执行官，于2011年接替史蒂夫·乔布斯。苹果公司在2023年发布了其首款混合现实头显设备Apple Vision Pro。 === 最终输出最优的匹配结果 === 问题：苹果公司的创始人是谁？答案：苹果公司于1976年由史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗纳德·韦恩创立，最初主要生产和销售个人电脑。</code>

<code class="javascript">import torchfrom transformers import AutoModelForSequenceClassification, AutoTokenizer class BGEReranker:    def __init__(self, model_name="BAAI/bge-reranker-base"):        self.tokenizer = AutoTokenizer.from_pretrained(model_name)        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)        self.model.eval()  # 设置为评估模式        def rerank(self, query: str, documents: list, top_k: int = 5):        """        对文档进行重排序                :param query: 用户查询        :param documents: 候选文档列表        :param top_k: 返回前K个文档        :return: 排序后的文档和分数        """        # 构建查询-文档对        pairs = [(query, doc) for doc in documents]                # 特征编码        with torch.no_grad():            inputs = self.tokenizer(                pairs,                 padding=True,                 truncation=True,                 return_tensors='pt',                 max_length=512            )                        # 计算相关性分数            scores = self.model(**inputs).logits.squeeze(dim=-1).float().numpy()                # 组合文档和分数        scored_docs = list(zip(documents, scores))                # 按分数降序排序        scored_docs.sort(key=lambda x: x[1], reverse=True)                return scored_docs[:top_k] # 使用示例if __name__ == "__main__":    # 初始化重排序模型    reranker = BGEReranker()        # 用户查询    query = "如何学习钢琴？"        # 初始检索结果（假设从FAISS获取）    initial_docs = [        "钢琴保养和清洁方法",        "音乐理论基础入门",        "钢琴购买指南：如何选择第一台钢琴",        "钢琴入门指法教程：从零开始学习",        "十大著名钢琴家及其作品",        "钢琴练习曲目推荐：适合初学者"    ]        print("初始检索结果:")    for i, doc in enumerate(initial_docs):        print(f"{i+1}. {doc}")        # 应用重排序    reranked_docs = reranker.rerank(query, initial_docs, top_k=3)        print("重排序后结果:")    for i, (doc, score) in enumerate(reranked_docs):        print(f"{i+1}. [分数: {score:.4f}] {doc}")     print("最优的匹配结果:")    print(reranked_docs[0][0])</code>

<code class="javascript">初始检索结果:1. 钢琴保养和清洁方法2. 音乐理论基础入门3. 钢琴购买指南：如何选择第一台钢琴4. 钢琴入门指法教程：从零开始学习5. 十大著名钢琴家及其作品6. 钢琴练习曲目推荐：适合初学者 重排序后结果:1. [分数: 1.7239] 钢琴入门指法教程：从零开始学习2. [分数: 0.7639] 音乐理论基础入门3. [分数: -1.1614] 钢琴练习曲目推荐：适合初学者 最优的匹配结果:钢琴入门指法教程：从零开始学习</code>