base_rag/examples/quick_start.py

83 lines
3.5 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):
# 使用带重排功能的搜索
print(f"查询参数 k={k}")
docs = self.similarity_search_with_rerank(question, k=k)
print(f"重排后返回文档数量: {len(docs)}")
return docs
def query_without_rerank(self, question, k=3):
# 不使用重排的普通搜索
docs = self.similarity_search(question, k=k)
return docs
def main():
# 重排配置 - 使用BGE重排模型
rerank_config = {
"enabled": True,
"type": "local",
"model": "BAAI/bge-reranker-base",
# 注意这里不设置top_k让它在查询时动态决定
}
rag = SimpleRAG(rerank_config=rerank_config)
print("RAG系统含重排功能初始化完成!")
# 添加更多测试文档
# documents = [
# "苹果是一种红色或绿色的水果味道甜美营养丰富含有丰富的维生素C。",
# "苹果公司是一家总部位于美国的科技公司以生产iPhone、iPad、Mac等产品而闻名。",
# "Python是一种高级编程语言简单易学功能强大广泛用于数据科学和机器学习。",
# "苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。",
# "苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。",
# "苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。",
# "iPhone是苹果公司生产的智能手机具有先进的技术和优秀的用户体验。",
# "机器学习是人工智能的一个分支Python是机器学习领域最流行的编程语言之一。",
# "苹果公司主要产品是iPhone、iPad、Mac等",
# "使用多线程可以提高程序在多核处理器上的运行效率。",
# "数据库设计应该遵循第三范式,以减少数据冗余。",
# "使用更高效的数据结构(如哈希表)可以提升算法效率。",
# "程序界面应该简洁明了,方便用户操作。",
# "尽量避免使用全局变量,减少潜在的并发问题。",
# "在Python中使用NumPy等库可以加快数值计算速度。",
# "使用SSD替代传统HDD可提升文件读写速度。",
# "将代码中的重复逻辑封装为函数,有助于维护但对性能影响小。",
# "内存泄漏会导致程序运行越来越慢,应及时释放资源。",
# "使用代码缓存如Memoization可以避免重复计算。",
# ]
# print("正在添加文档...")
# rag.ingest(documents)
# print(f"文档添加完成! 共添加了 {len(documents)} 个文档")
# 测试查询并比较重排效果
test_query = "Python中的list操作较慢尽量使用数组库"
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}")
if __name__ == "__main__":
main()