使用 LlamaIndex 实现 RAG

安装必要的库

1pip install llama-index llama-index-llms-groq ollama llama-index-embeddings-ollama

使用 Ollama 嵌入对文档进行索引

1from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
2from llama_index.embeddings.ollama import OllamaEmbedding
3
4# 从目录加载文档
5documents = SimpleDirectoryReader(input_dir="/content/data").load_data()
6
7# 定义嵌入模型
8embedding_model = OllamaEmbedding(model_name="nomic-embed-text")
9
10# 使用嵌入模型创建索引
11index = VectorStoreIndex.from_documents(documents, embed_model=embedding_model)
12index.storage_context.persist()

基于相似度分数检索文档

1# 第一步：从索引中检索相关节点
2retriever = index.as_retriever(similarity_top_k=2)
3retrieved_nodes = retriever.retrieve("谁提出了运动定律？")
4
5# 第二步：显示检索到的节点内容及其分数
6for i, node_with_score in enumerate(retrieved_nodes, start=1):
7    print(f"结果 {i}")
8    print("分数:", node_with_score.score)
9    print("内容:", node_with_score.node.get_content())
10    print("-" * 50)

结合检索与生成的查询引擎

1from llama_index.llms.groq import Groq
2
3# 加载一个大语言模型
4llm = Groq(model="llama3-70b-8192", api_key="YOUR_GROQ_API_KEY")
5
6# 创建一个集成了检索和响应生成功能的查询引擎
7query_engine = index.as_query_engine(llm=llm)
8response = query_engine.query("相对论是谁提出的？")
9
10print(response)