这是 LlamaIndex 官方 Starter Tutorial 中 demo,用很少的代码来使用 OpenAI 快速实现出一个 RAG。
Ref: Starter Tutorial | LlamaIndex
代码:llamindex-rag-demo | Kaggle
1)设置 OpenAI Token
这里使用国内的 OpenAI 中转 API token,在 LlamaIndex 中需要进行如下设置:
import os
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core.settings import SettingsOPENAI_API_KEY = 'sk-xxxx'
OPENAI_API_BASE = 'https://xxxx'os.environ['OPENAI_API_KEY'] = OPENAI_API_KEY
os.environ['OPENAI_API_BASE'] = OPENAI_API_BASESettings.llm = OpenAI(model="gpt-3.5-turbo", api_key=OPENAI_API_KEY, api_base=OPENAI_API_BASE)
Settings.embed_model = OpenAIEmbedding(api_key=OPENAI_API_KEY, api_base=OPENAI_API_BASE)
其中 Settings 为设置整个代码运行所使用的默认 LLM。
2)加载数据,并构建向量索引
这里使用一个 Paul Graham 的文章 “What I Worked On” 作为本地知识语料库的数据。
语料库下载连接:download
将改 txt 文件下载到 data
目录下。
从 data 目录加载语料库:
from llama_index.core import SimpleDirectoryReaderdocuments = SimpleDirectoryReader("data").load_data()
其中 documents
是一个 Document 对象的列表,每个 Document 对象包含了文件类型、文件长度、文件内容等各种文档描述信息:
基于 documents 可以构建向量索引,这里使用基于内存的向量索引 VectorStoreIndex
:
from llama_index.core import VectorStoreIndexindex = VectorStoreIndex.from_documents(documents)
3)使用 user query 进行 QA
这里从 index 构建出 Query Engine 用于做 QA。
官方文档对 Engine 的解释如下:
Engines provide natural language access to your data. For example:
- Query engines are powerful interfaces for question-answering (e.g. a RAG pipeline).
- Chat engines are conversational interfaces for multi-message, “back and forth” interactions with your data.
代码如下:
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(response)
response 就是检索增强后的生成响应:
以上就是一个简单的 RAG demo,具体如何打印日志、持久化 index 可以参考官方文档。