base_rag/examples/quick_start.py

92 lines
3.2 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 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(
embedding_config=embedding_config,
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()