这张照片由 Marc Sendra Martorell 在 Unsplash 上拍摄
- 开始
近期，检索增强生成（简称 RAG）已成为使用大型语言模型构建生成式 AI 应用程序的实际上的标准。RAG 增强了文本生成，通过确保生成模型使用适当的上下文，同时避免了微调 LLM 以执行相同任务所需的时间、成本和复杂性。RAG 还允许更高效地利用外部数据源，这使得模型的知识更新更加简便。
尽管基于RAG的AI应用通常使用更小或更精简的LLM，它们依然依赖于一个强大的管道，该管道能够嵌入和索引需要的知识库，并能高效检索并注入相关上下文到模型提示。
在许多应用场景中，RAG 可以利用任何优秀的框架在短短几行代码中实现。本文主要关注更复杂且要求更高的管道，例如，当需要嵌入和索引的数据量较大时，或者当数据需要频繁更新或快速更新时。
这篇文章演示了如何设计一个Rust应用程序，以惊人的速度读取、分块、嵌入并以向量形式存储文本文档。利用HuggingFace的Candle框架和LanceDB，它展示了如何构建一个端到端的RAG索引管道，可以部署在任何地方作为一个独立的应用程序，并且即使在非常严苛和孤立的环境中，也可以作为构建强大管道的基石。
本文的主要目的是创建一个可以应用于实际案例的工作示例，同时引导读者理解其关键设计原则和构建模块。该应用程序及其源代码可在随附的 GitHub 仓库（见下链接）中找到，可以直接使用，或作为进一步开发的参考。
本文的结构如下：第二部分（本节）概述了主要的设计决策和相关组件。第三部分详细介绍了管道的主要流程和组件的设计。第四部分和第五部分分别讨论了嵌入过程和写操作。第六部分是结论。
2. 设计选择和关键要素我们的主要设计目标是构建一个能够独立运行的应用程序，它可以运行一个端到端的索引处理流程，无需依赖外部服务或服务器进程。输出将是一组遵循LanceDB的Lance格式的数据文件，这些文件可以被LangChain、Llamaindex等框架使用，也可以通过DuckDB或其他使用LanceDB API的应用程序查询。
这个应用会用 Rust 编写，并基于两个主要的开源框架：将使用 Candle ML 框架来生成类似 BERT 模型的文档嵌入，并将使用 LanceDB 作为向量数据库和检索 API。
一个由作者制作的 Rust 应用程序，用于处理文档索引流程的所有阶段（图片由作者提供）
在详细了解应用程序的结构和细节之前，先简单聊聊这些组件和设计选择可能更有帮助。
在性能至关重要的情况下，Rust 是一个显而易见的选择。虽然 Rust 的学习曲线较陡，但其性能可与 C 或 C++ 等原生编程语言相媲美，并提供了丰富的抽象库和扩展功能，使得内存安全和并发等挑战比在原生语言中更容易处理。结合 Hugging Face 的 Candle 框架后，使用原生 Rust 中的 LLM 和嵌入模型从未如此顺畅。
然而，LanceDB 是最近才成为 RAG 技术堆栈的一员。它是一个轻量级且嵌入式的向量数据库，类似于 SQLite，可以直接集成到应用程序中，而无需单独的服务器进程。因此，它可以在任何地方部署，并嵌入到任何应用程序中，同时提供极快的搜索和检索能力，即使是对远端对象存储（例如 AWS S3）中的数据也是如此。正如我们之前提到的，它还提供了与 LangChain 和 LlamaIndex 的集成，并且可以通过 DuckDB 进行查询，这使得它成为一个更有吸引力的选择。
在我10核的Mac（没有使用GPU加速）上进行的一次简单测试中，该应用在一秒钟内处理、嵌入和存储了大约25,000个单词（相当于17个文本文件，每个文件约有1,500个单词）。这令人印象深刻的一次吞吐量展示了Rust在CPU密集型任务和I/O操作中的高效性，以及LanceDB强大的存储功能。这种组合在处理大规模数据嵌入和索引挑战方面表现出色。
这张照片由 Tharoushan Kandarajah 拍摄，来自 Unsplash
3. 管道架构与数据流我们的RAG应用和索引流程主要包括两个主要任务：“读取和嵌入任务”，该任务从文本文件中读取文本并使用嵌入模型将其嵌入到BERT向量中，以及“写入任务”，该任务将嵌入写入向量存储。由于前者主要受CPU限制（嵌入单个文档可能需要多次调用机器学习模型），而后者主要等待IO操作，我们将这两个任务分配到不同的线程。此外，为了防止争用和背压现象，我们将通过MPSC通道将这两个任务连接起来。在Rust（和其他语言中），同步通道基本上使线程间的安全异步通信成为可能，从而更好地扩展。
主要流程很简单：每次嵌入任务将一个文本文档嵌入成向量后，它会将该向量及其ID（即文件名）发送到通道，然后立即继续处理下一个文档（请参见下图中的读取部分）。同时，写入任务会不断从通道中读取数据，在内存中分块向量，并在达到一定大小时将其刷新至内存。因为预计嵌入任务会比较耗时且占用较多资源，我们将并行执行以充分利用运行应用程序的机器上的所有可用核心。换句话说，我们将有多个嵌入任务来读取并嵌入文档，而只有一个写入任务来分块并将向量写入数据库中。
管道设计及应用流程，（作者供图）
我们从 main() 函数开始吧，这会让整个流程更清晰易懂。
如上所示，在设置好通道（第3行）之后，我们初始化写任务线程，该线程开始从通道中获取消息，直到通道关闭。接下来，它会列出相关目录中的文件并将这些文件存储在一个字符串列表中。最后，它使用_Rayon_的par_iter函数来并行处理文件列表，以实现使用process_text_file()函数并行处理文件。使用Rayon可以利用机器的最大性能来最大化并行处理文档。
如您所见，流程相对简单明了，主要涉及两个核心任务：文档处理和向量存储。这种设计不仅使并行处理变得高效，还提升了系统的可扩展性。文档处理任务使用 Rayon 并行化文件处理，充分利用系统资源。同时，存储任务高效地将嵌入向量写入 LanceDB。这种职责分离不仅简化了整体架构，还使得每个任务可以独立优化。在接下来的内容中，我们将更深入地探讨这两个方面。
4 文档嵌入技术与Candle正如我们之前讨论的，在我们管道的一端，我们有一系列嵌入任务，每个任务都在独立的线程上运行。Rayon的iter_par函数有效遍历文件列表，对每个文件调用process_text_file()函数，同时最大化并行处理。
我们从这个函数开始吧
该函数首先获取自身的嵌入模型的引用（这是函数中最复杂的部分，我会稍后详细解释）。接着，它将文件分块读取，并对每个块调用嵌入函数（该函数基本上调用了模型本身）。嵌入函数返回一个 Vec<f32> 类型的向量（大小为 [1, 384]），这是对每个块进行嵌入和归一化处理后的结果，然后对所有文本块计算平均值。完成这一部分后，将向量和文件名一起发送到通道，用于持久化、查询以及由写入任务进行检索。
这里大部分的工作是由 BertModelWrapper 结构（在第 2 行中获取其引用）完成的。BertModelWrapper 的主要目的是封装模型的加载和嵌入过程，并提供 embed_sentences() 函数，该函数的主要功能是将一组文本片段嵌入，并计算它们的平均值来生成一个向量，从而生成一个单个向量。
为了实现该目标，BertModelWrapper 使用了 HuggingFace 的 Candle 框架。Candle 是一个原生的 Rust 库，它的 API 与 PyTorch 相似，用来加载和管理 ML 模型，并且在 HuggingFace 上托管的模型中，它提供了非常方便的支持。尽管在 Rust 中还有其他生成文本嵌入的方法，但 Candle 由于其原生性及不依赖其他库，在这方面看起来是最简洁的。
虽然对包装器代码的详细说明超出了我们当前的范围，我在另一篇文章中详细介绍了这个内容，并且其源代码也在附带的 GitHub 仓库中提供。您还可以在 Candle 的示例仓库中找到很多优秀的示例。
然而，有一个重要的部分需要解释我们是如何使用嵌入模型的，因为这将在整个流程中需要大规模使用模型时成为一个挑战。简而言之：我们希望我们的模型可以被多个线程用于执行嵌入任务，但因为它的加载时间较长，我们并不希望每次需要时都重新加载模型。换句话说，我们希望确保每个线程只会创建一个模型实例来处理嵌入任务，并重复使用该实例来为多个嵌入任务生成嵌入向量。
由于 Rust 的一些众所周知的限制，实现这些要求并不容易。如果您不想深入了解这些实现细节在 Rust 中是如何实现的，可以跳过这一部分（可以直接用代码）。
我们先从获取模型引用的这个函数开始：
我们的模型通过几个层次的包装来实现上述功能。首先，它被包裹在一个 thread_local 语句中，这意味着每个线程都有自己的懒加载副本，也就是说每个线程都可以访问 BERT_MODEL。但是只有当第一次调用 with()（第18行）时才会触发初始化代码，只会被懒加载，每个线程只会执行一次初始化，确保每个线程都有一个只初始化一次的有效引用。第二层是引用计数类型——Rc，这样就能更简便地创建模型的引用而无需考虑生命周期问题。每次调用 clone()，我们都会得到一个新的引用，这个引用会在超出其作用域时自动释放。
最后一层本质上就是服务功能 get_model_reference()，它简单地调用了 with() 函数，提供对内存中线程局部区域访问，该区域持有初始化的模型。调用 clone() 将给我们提供一个线程局部的模型引用，如果模型还没有初始化，那么初始化代码会被首先执行。
现在我们知道了如何同时运行多个嵌入任务，并且将向量写到通道中，我们可以继续到管道的另一部分——写任务。
写作任务相对简单一些，主要作为封装了 LanceDB 写入功能的接口层。回想一下，LanceDB 是一个作为库形式的查询引擎，可以读写存储在远程存储（如 AWS S3）上的数据，但并不拥有这些数据。这使得在这种场景中尤其方便，特别是在需要处理大规模数据时，可以保持低延迟性而无需管理单独的数据库服务器。
LanceDB 的 Rust API 使用 Arrow 来定义数据模型和表示数据，其 Python API 对某些人来说可能更方便。例如，这是我们用 Arrow 格式定义数据模型的方式：
如您所见，我们当前的架构包含两个字段：一个“filename”字段，保存实际文件位置，并用作我们的键，以及一个“vector”字段，保存实际文档的向量。在LanceDB中，向量使用FixedSizeList Arrow类型（表示一个数组）表示，而向量中的每个元素是Float32类型。（向量的长度为384，这是最后设置的。）
连接到LanceDB很简单，只需要提供一个存储位置，它可以是本地路径或S3 URI。然而，使用Rust和Arrow数据结构向LanceDB追加数据则不太友好。类似于其他基于Arrow的列式数据结构，追加操作不是将行列表添加进去，而是每个列都表示为一个值列表的方式进行。例如，如果你需要插入10行且包含2个列，你需要为每个列追加一个包含10个值的列表。
这里有一个例子，
代码的核心位于第2行，我们根据模式和列数据构建了一个Arrow RecordBatch。在这种情况下，我们有两个列——filename和vector。我们使用两个列表初始化记录批次：key_array，一个包含文件名的字符串列表，和vectors_array，一个包含向量的数组列表。从这里开始，由于Rust严格的类型安全要求，我们需要对这些数据进行大量的封装，之后才能传递给在第1行获取的表引用的add()函数。
为了简化这一逻辑，我们创建了一个存储模块来封装这些操作并，提供了基于connect(uri)和add_vector的简单接口。以下是写入任务线程的完整代码，该线程从通道读取嵌入，对其进行分块，并在达到一定大小时写入。
一旦数据被写入后，LanceDB的数据文件可以从任何进程访问。这里有一个例子，展示我们如何可以使用相同的数据进行向量相似度搜索，通过LanceDB的Python API，该API可以在完全不同的进程中运行。
uri = "data/vecdb1"
db = lancedb.connect(uri)
tbl = db.open_table("vectors_table_1")
# 我们要找的相似向量
encoded_vec = get_some_vector()
# 搜索与之最相似的前3个向量
tbl.search(embeddings[0]) \
.select(["filename"]) \
.limit(3).to_pandas()
脚本生成的图片
5. 结论在这篇文章中，我们看到了一个使用Rust、HuggingFace的Candle框架和LanceDB构建的高性能RAG管道流程的工作示例。我们了解了如何利用Rust的性能优势，结合Candle框架，高效并行读取和嵌入多个文本文件。我们还看到了如何使用同步通道并行运行嵌入任务和写入流程，而不需要处理复杂的锁定和同步问题。最后，我们还了解到如何利用LanceDB高效的存储特性，用Rust生成可以与多个AI框架和查询库无缝对接的向量数据库。
我认为这里描述的方法可以作为一个构建可扩展、生产就绪的检索增援生成（RAG）索引管道的强大基础。无论你是处理大量数据，需要频繁更新知识库，还是在资源受限的环境中工作，文中讨论的构建模块和设计原则都可以根据你的具体需求进行调整。随着人工智能领域的不断进步，高效处理和检索相关信息的能力将一直保持重要性。如本文所展示，结合合适的工具和周到的设计，开发人员可以构建不仅满足当前需求，还能够应对未来人工智能驱动的信息检索和生成挑战的RAG管道。
注释与链接
- 您可以在此处找到附带源代码的 GitHub 仓库此处。此外，该仓库还包含一个示例 Jupyter 笔记本，演示了如何使用 Python 进行测试。
- 我关于 HuggingFace Candle 的上一篇文章可以在这里找到此处。
- Candle 框架及其文档，包括其详尽的示例文件夹
- LanceDB及其Rust API 文档，查看详细信息。
随时随地看视频
