base_rag/examples/quick_start.py

47 lines
1.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 SimpleRAG(BaseRAG):
def ingest(self, documents):
for doc in documents:
self.vector_store.add_texts([doc])
def query(self, question, k=3):
docs = self.vector_store.similarity_search(question, k=k)
return docs
def main():
config = {
"model_name": "sentence-transformers/all-MiniLM-L6-v2",
"embedding_type": "local",
}
rag = SimpleRAG(embedding_config=config)
print("RAG初始化完成!")
# 添加一些文档
documents = [
"苹果是一种水果,味道甜美,营养丰富。",
"苹果公司是一家科技公司生产iPhone和Mac等产品。",
"Python是一种编程语言简单易学功能强大。",
]
print("正在添加文档...")
rag.ingest(documents)
print("文档添加完成!")
# 测试查询
print("\n正在查询: '什么是苹果?'")
result = rag.query("什么是苹果?")
print(f"查询结果: {result}")
if __name__ == "__main__":
main()