diff --git a/examples/quick_start.py b/examples/quick_start.py index 6d21c5c..522c1e4 100644 --- a/examples/quick_start.py +++ b/examples/quick_start.py @@ -13,7 +13,9 @@ class SimpleRAG(BaseRAG): def query(self, question, k=3): # 使用带重排功能的搜索 + print(f"查询参数 k={k}") docs = self.similarity_search_with_rerank(question, k=k) + print(f"重排后返回文档数量: {len(docs)}") return docs def query_without_rerank(self, question, k=3): @@ -23,34 +25,46 @@ class SimpleRAG(BaseRAG): def main(): - # 重排配置 - 使用基于余弦相似度的重排 + # 重排配置 - 使用BGE重排模型 rerank_config = { "enabled": True, "type": "local", "model": "BAAI/bge-reranker-base", - "top_k": 3, + # 注意:这里不设置top_k,让它在查询时动态决定 } rag = SimpleRAG(rerank_config=rerank_config) print("RAG系统(含重排功能)初始化完成!") # 添加更多测试文档 - documents = [ - "苹果是一种红色或绿色的水果,味道甜美,营养丰富,含有丰富的维生素C。", - "苹果公司是一家总部位于美国的科技公司,以生产iPhone、iPad、Mac等产品而闻名。", - "Python是一种高级编程语言,简单易学,功能强大,广泛用于数据科学和机器学习。", - "苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。", - "苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。", - "苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。", - "iPhone是苹果公司生产的智能手机,具有先进的技术和优秀的用户体验。", - "机器学习是人工智能的一个分支,Python是机器学习领域最流行的编程语言之一。", - ] - print("正在添加文档...") - rag.ingest(documents) - print(f"文档添加完成! 共添加了 {len(documents)} 个文档") + # documents = [ + # "苹果是一种红色或绿色的水果,味道甜美,营养丰富,含有丰富的维生素C。", + # "苹果公司是一家总部位于美国的科技公司,以生产iPhone、iPad、Mac等产品而闻名。", + # "Python是一种高级编程语言,简单易学,功能强大,广泛用于数据科学和机器学习。", + # "苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。", + # "苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。", + # "苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。", + # "iPhone是苹果公司生产的智能手机,具有先进的技术和优秀的用户体验。", + # "机器学习是人工智能的一个分支,Python是机器学习领域最流行的编程语言之一。", + # "苹果公司主要产品是iPhone、iPad、Mac等", + # "使用多线程可以提高程序在多核处理器上的运行效率。", + # "数据库设计应该遵循第三范式,以减少数据冗余。", + # "使用更高效的数据结构(如哈希表)可以提升算法效率。", + # "程序界面应该简洁明了,方便用户操作。", + # "尽量避免使用全局变量,减少潜在的并发问题。", + # "在Python中使用NumPy等库可以加快数值计算速度。", + # "使用SSD替代传统HDD可提升文件读写速度。", + # "将代码中的重复逻辑封装为函数,有助于维护但对性能影响小。", + # "内存泄漏会导致程序运行越来越慢,应及时释放资源。", + # "使用代码缓存(如Memoization)可以避免重复计算。", + + # ] + # print("正在添加文档...") + # rag.ingest(documents) + # print(f"文档添加完成! 共添加了 {len(documents)} 个文档") # 测试查询并比较重排效果 - test_query = "苹果的营养价值和健康效益" + test_query = "Python中的list操作较慢,尽量使用数组库" print(f"\n测试查询: '{test_query}'") print("\n=== 不使用重排的结果 ===") @@ -63,20 +77,6 @@ def main(): for i, doc in enumerate(result_with_rerank, 1): print(f"{i}. {doc.page_content}") - # 另一个测试查询 - test_query2 = "苹果公司的主要产品" - print(f"\n\n测试查询2: '{test_query2}'") - - print("\n=== 不使用重排的结果 ===") - result_no_rerank2 = rag.query_without_rerank(test_query2, k=3) - for i, doc in enumerate(result_no_rerank2, 1): - print(f"{i}. {doc.page_content}") - - print("\n=== 使用重排的结果 ===") - result_with_rerank2 = rag.query(test_query2, k=3) - for i, doc in enumerate(result_with_rerank2, 1): - print(f"{i}. {doc.page_content}") - if __name__ == "__main__": main() diff --git a/src/base_rag/core.py b/src/base_rag/core.py index 0fc95b1..7256bed 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -231,13 +231,17 @@ class BaseRAG(ABC): ) def _rerank_documents( - self, query: str, documents: List[Document] + self, query: str, documents: List[Document], top_k: int = None ) -> List[Document]: """对检索到的文档进行重排""" if not documents: return documents - top_k = self.rerank_config.get("top_k", len(documents)) + # 优先使用传入的top_k,否则使用配置中的top_k,最后使用文档总数 + if top_k is None: + top_k = self.rerank_config.get("top_k", len(documents)) + + print(f"重排输入: {len(documents)} 个文档, 目标返回: {top_k} 个") # 检查reranker模型是否可用 if not self.reranker: @@ -253,16 +257,25 @@ class BaseRAG(ABC): query_doc_pairs = [(query, doc.page_content) for doc in documents] scores = self.reranker.predict(query_doc_pairs) - # 处理单个文档的情况 - if not isinstance(scores, list): + # 处理得分数据:确保scores是一维列表 + if isinstance(scores, np.ndarray): + scores = scores.flatten().tolist() + elif not isinstance(scores, list): scores = [scores] + elif len(scores) == 1 and isinstance(scores[0], np.ndarray): + # 如果是包含单个数组的列表,提取数组内容 + scores = scores[0].flatten().tolist() + + print(f"重排得分: {scores}") # 根据分数排序 doc_scores = list(zip(documents, scores)) doc_scores.sort(key=lambda x: x[1], reverse=True) - + # 返回top_k个文档 - return [doc for doc, score in doc_scores[:top_k]] + result = [doc for doc, score in doc_scores[:top_k]] + print(f"重排输出: {len(result)} 个文档") + return result except Exception as e: print(f"重排失败: {e},跳过重排") @@ -377,11 +390,13 @@ class BaseRAG(ABC): # 如果启用了重排,进行重排 if self.rerank_config.get("enabled", False) and documents: - documents = self._rerank_documents(query, documents) - - # 返回最终的top_k结果 - final_k = k or self.retriever_top_k - return documents[:final_k] + final_k = k or self.retriever_top_k + documents = self._rerank_documents(query, documents, top_k=final_k) + return documents + else: + # 返回最终的top_k结果 + final_k = k or self.retriever_top_k + return documents[:final_k] @abstractmethod def ingest(self, *args, **kwargs):