140 lines
5.2 KiB
Python
140 lines
5.2 KiB
Python
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": "similarity",
|
||
"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()
|