90 lines
3.2 KiB
Python
90 lines
3.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 SimpleRAG(BaseRAG):
|
||
def ingest(self, documents):
|
||
for doc in documents:
|
||
self.vector_store.add_texts([doc])
|
||
|
||
def query(self, question, k=3):
|
||
# 使用带重排功能的搜索
|
||
docs = self.similarity_search_with_rerank(question, k=k)
|
||
return docs
|
||
|
||
def query_without_rerank(self, question, k=3):
|
||
# 不使用重排的普通搜索
|
||
docs = self.similarity_search(question, k=k)
|
||
return docs
|
||
|
||
|
||
def main():
|
||
# 嵌入模型配置
|
||
embedding_config = {
|
||
"type": "local",
|
||
"model_name": "sentence-transformers/all-MiniLM-L6-v2",
|
||
}
|
||
|
||
# 重排配置 - 使用基于余弦相似度的重排
|
||
rerank_config = {
|
||
"enabled": True,
|
||
"method": "similarity", # 使用相似度重排(无需额外依赖)
|
||
"top_k": 3
|
||
}
|
||
|
||
rag = SimpleRAG(
|
||
rerank_config=rerank_config
|
||
)
|
||
print("RAG系统(含重排功能)初始化完成!")
|
||
|
||
# 添加更多测试文档
|
||
documents = [
|
||
"苹果是一种红色或绿色的水果,味道甜美,营养丰富,含有丰富的维生素C。",
|
||
"苹果公司是一家总部位于美国的科技公司,以生产iPhone、iPad、Mac等产品而闻名。",
|
||
"Python是一种高级编程语言,简单易学,功能强大,广泛用于数据科学和机器学习。",
|
||
"苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。",
|
||
"苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。",
|
||
"苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。",
|
||
"iPhone是苹果公司生产的智能手机,具有先进的技术和优秀的用户体验。",
|
||
"机器学习是人工智能的一个分支,Python是机器学习领域最流行的编程语言之一。"
|
||
]
|
||
print("正在添加文档...")
|
||
rag.ingest(documents)
|
||
print(f"文档添加完成! 共添加了 {len(documents)} 个文档")
|
||
|
||
# 测试查询并比较重排效果
|
||
test_query = "苹果的营养价值和健康效益"
|
||
print(f"\n测试查询: '{test_query}'")
|
||
|
||
print("\n=== 不使用重排的结果 ===")
|
||
result_no_rerank = rag.query_without_rerank(test_query, k=3)
|
||
for i, doc in enumerate(result_no_rerank, 1):
|
||
print(f"{i}. {doc.page_content}")
|
||
|
||
print("\n=== 使用重排的结果 ===")
|
||
result_with_rerank = rag.query(test_query, k=3)
|
||
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()
|