使用 LlamaIndex 构建 RAG 电子书“图书管理员”

引言

本笔记本演示如何快速为本地电子书库构建一个基于 RAG 的“图书管理员”。回想一下你上次去图书馆时，是不是曾借助馆员丰富的知识，从大量教科书、小说和其他资源中找到所需资料？我们的 RAG“图书管理员”也会这样帮助我们，只不过它服务的是我们自己的本地电子书收藏。

需求

我们希望这个图书管理员足够轻量，尽可能在本地运行，并且依赖项尽量少。因此，我们会充分利用开源工具，优先选择能在普通硬件上本地运行的模型，例如 M1 MacBook。

组件

我们的方案包含以下组件：

• LlamaIndex：一个面向 LLM 应用的数据框架；与 LangChain 不同，它专门为 RAG 设计。
• Ollama：一个便于在本地运行 Llama 2 等 LLM 的工具。
• BAAI/bge-base-en-v1.5 嵌入模型：效果不错，体积也相对轻量。
• Llama 2：我们将通过 Ollama 运行它。

依赖项

首先安装依赖项。

%pip install -q \
    llama-index \
    EbookLib \
    html2text \
    llama-index-embeddings-huggingface \
    llama-index-llms-ollama

安装 Ollama

以下依赖项帮助正确检测 GPU。

!apt install pciutils lshw

安装 Ollama。

!curl -fsSL https://ollama.com/install.sh | sh

在后台运行 Ollama 服务。

get_ipython().system_raw("ollama serve &")

从 Ollama 模型库拉取 Llama 2。

!ollama pull llama2

准备测试书库

接下来创建测试“书库”。为简单起见，假设我们的“书库”就是一个嵌套的 .epub 文件目录。很容易想到，这个方案可以扩展到 Calibre 书库等场景，其中会有一个 metadata.db 数据库文件。我们将这项扩展留作读者练习。😇

从 Project Gutenberg 下载两个 .epub 文件作为我们的书库。

!mkdir -p "./test/library/jane-austen"
!mkdir -p "./test/library/victor-hugo"
!wget https://www.gutenberg.org/ebooks/1342.epub.noimages -O "./test/library/jane-austen/pride-and-prejudice.epub"
!wget https://www.gutenberg.org/ebooks/135.epub.noimages -O "./test/library/victor-hugo/les-miserables.epub"

使用 LlamaIndex 实现 RAG

从核心流程看，使用 LlamaIndex 实现 RAG 大致包含以下阶段：

1. 加载：告诉 LlamaIndex 数据存放在哪里，以及如何加载。
2. 索引：为已加载的数据添加向量嵌入等信息，以便进行查询。
3. 查询：配置一个 LLM，作为已建立索引的数据的查询接口。

这段说明只是触及 LlamaIndex 能力的表面。若想深入了解，强烈建议阅读 LlamaIndex 文档中的“高层概念”页面。

加载

自然，我们先从加载阶段开始。

前面提到，LlamaIndex 是专门为 RAG 设计的。这一点从它的 SimpleDirectoryReader 构造器中就能立即看出来：它开箱即用，✨ 神奇地 ✨支持一大批不同格式的文件。对我们来说很方便的是，.epub 也在支持范围内。

from llama_index.core import SimpleDirectoryReader

loader = SimpleDirectoryReader(
    input_dir="./test/",
    recursive=True,
    required_exts=[".epub"],
)

documents = loader.load_data()

SimpleDirectoryReader.load_data() 会将电子书转换为一组 Document，供 LlamaIndex 使用。这里有一点需要注意：此时文档尚未分块，分块会在索引阶段进行。继续往下看……

索引

加载数据后，下一步是建立索引。这样，RAG 流水线就能查找与查询相关的上下文，并将这些上下文传给 LLM，用来增强生成的回答。文档分块也会在这里发生。

VectorStoreIndex 是 LlamaIndex 中建立索引的“默认”入口。默认情况下，VectorStoreIndex 使用一个简单的内存字典保存索引；但 LlamaIndex 也支持多种向量存储方案，方便你在规模扩大时升级。

LlamaIndex 默认的分块大小为 1024，块间重叠为 20。详情参见 LlamaIndex 文档。

如前所述，我们将使用 BAAI/bge-small-en-v1.5 生成嵌入。默认情况下，LlamaIndex 使用 OpenAI（具体为 gpt-3.5-turbo）；我们希望获得轻量、可在本地端到端运行的方案，因此想避开这个默认行为。好在 LlamaIndex 支持通过方便的 HuggingFaceEmbedding 类从 Hugging Face 获取嵌入模型，我们就在这里使用它。

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embedding_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

将该模型作为嵌入模型传入 VectorStoreIndex，即可绕过默认的 OpenAI 行为。

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(
    documents,
    embed_model=embedding_model,
)

查询

现在补上 RAG 拼图的最后一块：接入查询层。本示例使用 Llama 2，但鼓励读者尝试不同模型，看看哪种模型能给出“最好”的回答。

首先启动 Ollama 服务器。遗憾的是，Ollama Python 客户端本身不支持启动或停止服务器，因此我们需要暂时离开 Python 环境。在另一个终端中运行 ollama serve。完成这里的操作后，记得终止该进程！

现在将 Llama 2 接入 LlamaIndex，并以它为基础建立查询引擎。

from llama_index.llms.ollama import Ollama

llama = Ollama(
    model="llama2",
    request_timeout=40.0,
)

query_engine = index.as_query_engine(llm=llama)

最终结果

这样，基本的 RAG 图书管理员就搭建好了，我们可以开始询问书库中的内容。例如：

>>> print(
...     query_engine.query(
...         "What are the titles of all the books available? Show me the context used to derive your answer."
...     )
... )

Based on the context provided, there are two books available:

1. "Pride and Prejudice" by Jane Austen
2. "Les Misérables" by Victor Hugo

The context used to derive this answer includes:

* The file path for each book, which provides information about the location of the book files on the computer.
* The titles of the books, which are mentioned in the context as being available for reading.
* A list of words associated with each book, such as "epub" and "notebooks", which provide additional information about the format and storage location of each book.

上述输出表示：根据给定上下文，有简·奥斯汀的《傲慢与偏见》和维克多·雨果的《悲惨世界》两本书。回答所使用的上下文包括每本书的文件路径、书名，以及 epub、notebooks 等与书籍关联的词，这些词提供了文件格式和存储位置的额外信息。

>>> print(query_engine.query("Who is the main character of 'Pride and Prejudice'?"))

The main character of 'Pride and Prejudice' is Elizabeth Bennet.

上述回答指出，《傲慢与偏见》的主角是伊丽莎白·班纳特。

结论与后续改进

我们演示了如何构建一个完全在本地运行的基本 RAG“图书管理员”，即使是在搭载 Apple 芯片的 Mac 上也可以运行。同时，我们也对 LlamaIndex 做了一次“总体巡览”，了解了它如何简化 RAG 应用的搭建过程。

不过，我们确实只触及了这里潜力的表面。以下是一些在这个基础上完善和扩展的思路。

强制提供引用

为了防范图书管理员产生幻觉的风险，我们如何要求它对自己说的每件事都提供引用来源？

使用扩展元数据

Calibre 等电子书库管理工具会为书库中的电子书创建额外的元数据。这些元数据能提供出版社、版本等信息，而这些信息未必容易从书籍正文中获得。我们如何扩展 RAG 流水线，使其利用 .epub 文件以外的信息来源？

高效建立索引

如果把这里构建的内容全部整理为一个脚本或可执行文件，那么每次调用它时都会重新索引书库。对于只有两个文件的微型测试书库，这“没问题”；但只要书库有一定规模，用户很快就会觉得烦人。

我们如何持久化嵌入索引，并且只在书库内容发生有意义的变化时（例如加入新书）更新索引？