import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from base_rag import BaseRAG class RerankRAG(BaseRAG): def __init__(self, **kwargs): super().__init__(**kwargs) self.document_count = 0 def ingest(self, documents): """批量添加文档,避免重复""" if documents: # 清空现有集合并重新添加所有文档 self.vector_store.delete_collection() # 重新初始化向量库 from langchain_chroma import Chroma self.vector_store = Chroma( collection_name=self.vector_store_name, embedding_function=self.embedding_model, persist_directory=self.persist_directory, ) # 添加所有文档 self.vector_store.add_texts(documents) self.document_count = len(documents) print(f"已添加 {self.document_count} 个文档到向量库") def query_with_scores(self, question, k=5): """带分数的查询,用于比较重排效果""" # 不使用重排的结果 docs_no_rerank = self.similarity_search(question, k=k) # 使用重排的结果 docs_with_rerank = self.similarity_search_with_rerank(question, k=k) return docs_no_rerank, docs_with_rerank def query(self, question, k=3): return self.similarity_search_with_rerank(question, k=k) def main(): # 嵌入模型配置 embedding_config = { "type": "local", "model_name": "sentence-transformers/all-MiniLM-L6-v2", } # 重排配置 rerank_config = {"enabled": True, "method": "cross_encoder", "top_k": 3} # 初始化RAG系统 rag = RerankRAG( vector_store_name="rerank_test", embedding_config=embedding_config, rerank_config=rerank_config, retriever_top_k=5 # 获取更多候选文档 ) print("RAG系统(含重排功能)初始化完成!\n") # 测试文档 - 关于不同主题的文档 documents = [ # 水果相关 "苹果是一种非常有营养的水果,富含维生素C、纤维和抗氧化剂,对心脏健康有益。", "橙子含有丰富的维生素C,是柑橘类水果的代表,有助于增强免疫系统。", "香蕉富含钾元素,能够帮助维持血压稳定,是运动员的理想能量补充。", # 科技公司相关 "苹果公司(Apple Inc.)是全球知名的科技公司,主要产品包括iPhone、iPad、Mac电脑等。", "谷歌公司专注于搜索引擎和云计算服务,Android操作系统是其重要产品。", "微软公司开发Windows操作系统和Office办公软件,在企业软件领域占据重要地位。", # 编程语言相关 "Python是一种高级编程语言,语法简洁,广泛用于数据科学、机器学习和Web开发。", "Java是面向对象的编程语言,具有跨平台特性,在企业级开发中应用广泛。", "JavaScript是Web开发的核心语言,可以实现网页的交互功能和动态效果。", # 健康相关 "规律运动有助于维持身体健康,建议每周至少进行150分钟的中等强度有氧运动。", "均衡饮食是健康的基础,应该多吃蔬菜水果,减少加工食品的摄入。", "充足的睡眠对身心健康至关重要,成年人每天应保证7-9小时的睡眠时间。" ] print("正在添加测试文档...") rag.ingest(documents) print(f"文档添加完成!\n") # 测试查询 test_queries = [ { "query": "苹果对健康有什么好处?", "expected_topic": "应该更偏向水果营养相关的文档" }, { "query": "苹果公司的主要业务是什么?", "expected_topic": "应该更偏向科技公司相关的文档" }, { "query": "如何保持身体健康?", "expected_topic": "应该更偏向健康建议相关的文档" } ] for i, test_case in enumerate(test_queries, 1): query = test_case["query"] expected = test_case["expected_topic"] print(f"=== 测试查询 {i}: {query} ===") print(f"期望结果: {expected}\n") # 获取两种搜索结果 docs_no_rerank, docs_with_rerank = rag.query_with_scores(query, k=5) print("📍 不使用重排的结果:") for j, doc in enumerate(docs_no_rerank[:3], 1): print(f" {j}. {doc.page_content}") print("\n🎯 使用重排的结果:") for j, doc in enumerate(docs_with_rerank[:3], 1): print(f" {j}. {doc.page_content}") print("\n" + "="*80 + "\n") print("重排功能测试完成!") print("\n说明:") print("- 重排功能通过计算查询与文档的余弦相似度来重新排序检索结果") print("- 理论上重排后的结果应该更符合查询意图") print("- 如果结果相同,可能是因为初始检索结果已经很好,或需要更多样化的测试数据") if __name__ == "__main__": main()