base_rag/examples/rerank_demo.py

136 lines
5.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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()