feat: 检测重排
This commit is contained in:
parent
1d7dd1f03b
commit
148474de9b
|
@ -13,7 +13,9 @@ class SimpleRAG(BaseRAG):
|
||||||
|
|
||||||
def query(self, question, k=3):
|
def query(self, question, k=3):
|
||||||
# 使用带重排功能的搜索
|
# 使用带重排功能的搜索
|
||||||
|
print(f"查询参数 k={k}")
|
||||||
docs = self.similarity_search_with_rerank(question, k=k)
|
docs = self.similarity_search_with_rerank(question, k=k)
|
||||||
|
print(f"重排后返回文档数量: {len(docs)}")
|
||||||
return docs
|
return docs
|
||||||
|
|
||||||
def query_without_rerank(self, question, k=3):
|
def query_without_rerank(self, question, k=3):
|
||||||
|
@ -23,34 +25,46 @@ class SimpleRAG(BaseRAG):
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
# 重排配置 - 使用基于余弦相似度的重排
|
# 重排配置 - 使用BGE重排模型
|
||||||
rerank_config = {
|
rerank_config = {
|
||||||
"enabled": True,
|
"enabled": True,
|
||||||
"type": "local",
|
"type": "local",
|
||||||
"model": "BAAI/bge-reranker-base",
|
"model": "BAAI/bge-reranker-base",
|
||||||
"top_k": 3,
|
# 注意:这里不设置top_k,让它在查询时动态决定
|
||||||
}
|
}
|
||||||
|
|
||||||
rag = SimpleRAG(rerank_config=rerank_config)
|
rag = SimpleRAG(rerank_config=rerank_config)
|
||||||
print("RAG系统(含重排功能)初始化完成!")
|
print("RAG系统(含重排功能)初始化完成!")
|
||||||
|
|
||||||
# 添加更多测试文档
|
# 添加更多测试文档
|
||||||
documents = [
|
# documents = [
|
||||||
"苹果是一种红色或绿色的水果,味道甜美,营养丰富,含有丰富的维生素C。",
|
# "苹果是一种红色或绿色的水果,味道甜美,营养丰富,含有丰富的维生素C。",
|
||||||
"苹果公司是一家总部位于美国的科技公司,以生产iPhone、iPad、Mac等产品而闻名。",
|
# "苹果公司是一家总部位于美国的科技公司,以生产iPhone、iPad、Mac等产品而闻名。",
|
||||||
"Python是一种高级编程语言,简单易学,功能强大,广泛用于数据科学和机器学习。",
|
# "Python是一种高级编程语言,简单易学,功能强大,广泛用于数据科学和机器学习。",
|
||||||
"苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。",
|
# "苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。",
|
||||||
"苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。",
|
# "苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。",
|
||||||
"苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。",
|
# "苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。",
|
||||||
"iPhone是苹果公司生产的智能手机,具有先进的技术和优秀的用户体验。",
|
# "iPhone是苹果公司生产的智能手机,具有先进的技术和优秀的用户体验。",
|
||||||
"机器学习是人工智能的一个分支,Python是机器学习领域最流行的编程语言之一。",
|
# "机器学习是人工智能的一个分支,Python是机器学习领域最流行的编程语言之一。",
|
||||||
]
|
# "苹果公司主要产品是iPhone、iPad、Mac等",
|
||||||
print("正在添加文档...")
|
# "使用多线程可以提高程序在多核处理器上的运行效率。",
|
||||||
rag.ingest(documents)
|
# "数据库设计应该遵循第三范式,以减少数据冗余。",
|
||||||
print(f"文档添加完成! 共添加了 {len(documents)} 个文档")
|
# "使用更高效的数据结构(如哈希表)可以提升算法效率。",
|
||||||
|
# "程序界面应该简洁明了,方便用户操作。",
|
||||||
|
# "尽量避免使用全局变量,减少潜在的并发问题。",
|
||||||
|
# "在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(f"\n测试查询: '{test_query}'")
|
||||||
|
|
||||||
print("\n=== 不使用重排的结果 ===")
|
print("\n=== 不使用重排的结果 ===")
|
||||||
|
@ -63,20 +77,6 @@ def main():
|
||||||
for i, doc in enumerate(result_with_rerank, 1):
|
for i, doc in enumerate(result_with_rerank, 1):
|
||||||
print(f"{i}. {doc.page_content}")
|
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__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
|
|
@ -231,13 +231,17 @@ class BaseRAG(ABC):
|
||||||
)
|
)
|
||||||
|
|
||||||
def _rerank_documents(
|
def _rerank_documents(
|
||||||
self, query: str, documents: List[Document]
|
self, query: str, documents: List[Document], top_k: int = None
|
||||||
) -> List[Document]:
|
) -> List[Document]:
|
||||||
"""对检索到的文档进行重排"""
|
"""对检索到的文档进行重排"""
|
||||||
if not documents:
|
if not documents:
|
||||||
return 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模型是否可用
|
# 检查reranker模型是否可用
|
||||||
if not self.reranker:
|
if not self.reranker:
|
||||||
|
@ -253,16 +257,25 @@ class BaseRAG(ABC):
|
||||||
query_doc_pairs = [(query, doc.page_content) for doc in documents]
|
query_doc_pairs = [(query, doc.page_content) for doc in documents]
|
||||||
scores = self.reranker.predict(query_doc_pairs)
|
scores = self.reranker.predict(query_doc_pairs)
|
||||||
|
|
||||||
# 处理单个文档的情况
|
# 处理得分数据:确保scores是一维列表
|
||||||
if not isinstance(scores, list):
|
if isinstance(scores, np.ndarray):
|
||||||
|
scores = scores.flatten().tolist()
|
||||||
|
elif not isinstance(scores, list):
|
||||||
scores = [scores]
|
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 = list(zip(documents, scores))
|
||||||
doc_scores.sort(key=lambda x: x[1], reverse=True)
|
doc_scores.sort(key=lambda x: x[1], reverse=True)
|
||||||
|
|
||||||
# 返回top_k个文档
|
# 返回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:
|
except Exception as e:
|
||||||
print(f"重排失败: {e},跳过重排")
|
print(f"重排失败: {e},跳过重排")
|
||||||
|
@ -377,11 +390,13 @@ class BaseRAG(ABC):
|
||||||
|
|
||||||
# 如果启用了重排,进行重排
|
# 如果启用了重排,进行重排
|
||||||
if self.rerank_config.get("enabled", False) and documents:
|
if self.rerank_config.get("enabled", False) and documents:
|
||||||
documents = self._rerank_documents(query, documents)
|
final_k = k or self.retriever_top_k
|
||||||
|
documents = self._rerank_documents(query, documents, top_k=final_k)
|
||||||
# 返回最终的top_k结果
|
return documents
|
||||||
final_k = k or self.retriever_top_k
|
else:
|
||||||
return documents[:final_k]
|
# 返回最终的top_k结果
|
||||||
|
final_k = k or self.retriever_top_k
|
||||||
|
return documents[:final_k]
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def ingest(self, *args, **kwargs):
|
def ingest(self, *args, **kwargs):
|
||||||
|
|
Loading…
Reference in New Issue