原文towardsdatascience.com/scale-up-your-rag-a-rust-powered-indexing-pipeline-with-lancedb-and-candle-cc681c6162e8?sourcecollection_archive---------2-----------------------#2024-07-11构建大规模文档处理的高性能嵌入和索引系统https://medium.com/alon.agmon?sourcepost_page---byline--cc681c6162e8--------------------------------https://towardsdatascience.com/?sourcepost_page---byline--cc681c6162e8-------------------------------- Alon Agmon·发表于Towards Data Science ·12 分钟阅读·2024 年 7 月 11 日–https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/96228249c0ebdc4684b0beec98e23e54.png图片来源Marc Sendra Martorell 来自Unsplash1. 简介最近基于检索增强生成RAG的 AI 应用已成为构建生成型 AI 应用程序的事实上的标准尤其是在使用大型语言模型时。RAG 通过确保生成模型使用适当的上下文来增强文本生成同时避免了对 LLM 进行微调所涉及的时间、成本和复杂性。RAG 还可以更高效地使用外部数据源并更容易地更新模型的“知识”。尽管基于 RAG 的 AI 应用通常可以使用更小型或更简单的 LLM但它们仍然依赖于一个强大的流程来嵌入和索引所需的知识库并且需要能够高效地检索并将相关上下文注入到模型提示中。在许多应用场景中RAG 可以通过使用任何一个广泛可用的优秀框架在几行代码中实现。本文侧重于更复杂且要求更高的流程例如当需要嵌入和索引的数据量较大或者需要非常频繁或极快地更新时。本文展示了如何设计一个 Rust 应用程序能够以惊人的速度读取、分块、嵌入并将文本文档存储为向量。利用 Hugging Face 的 Candle 框架和 LanceDB它展示了如何开发一个端到端的 RAG 索引管道可以作为独立应用程序部署到任何地方并作为强大管道的基础即使在非常苛刻和孤立的环境中也是如此。本文的主要目的是创建一个可以应用于现实世界用例的工作示例同时引导读者了解其关键设计原则和构建模块。该应用程序及其源代码可在随附的 GitHub 仓库中获得链接见下可以直接使用或作为进一步开发的示例。本文的结构如下第二部分高层次地解释了主要的设计选择和相关组件。第三部分详细介绍了管道的主要流程和组件设计。第四部分和第五部分分别讨论了嵌入流程和写入任务。第六部分作结。2. 设计选择与关键组件我们的主要设计目标是构建一个独立的应用程序能够在没有外部服务或服务器进程的情况下运行端到端的索引管道。其输出将是一个数据文件集采用 LanceDB 的Lance 格式这些文件可以被像 LangChain 或 Llamaindex 这样的框架使用并且可以通过 DuckDB 或任何使用 LanceDB API 的应用程序进行查询。该应用程序将用 Rust 编写并基于两个主要的开源框架我们将使用Candle ML框架处理生成文档嵌入的机器学习任务采用类似 BERT 的模型并使用LanceDB作为我们的向量数据库和检索 API。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/8fc3ab1d758e41063ba4ab6113b0840a.png处理文档索引管道所有阶段的 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 强大的存储能力。两者结合在一起对于解决大规模数据嵌入和索引挑战表现出色。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/038d06947393f658500ad6d318aa8974.png图片来源Tharoushan Kandarajah 在Unsplash上的作品3. 流水线架构与流程我们的 RAG 应用程序和索引流水线包含两个主要任务一个读取和嵌入任务它从文本文件中读取文本并使用嵌入模型将其嵌入到 BERT 向量中另一个是写入任务它将嵌入写入向量存储。由于前者大多由 CPU 限制嵌入单个文档可能需要多个机器学习模型操作而后者大多是在等待 I/O因此我们将这两个任务分配到不同的线程中。此外为了避免争用和背压我们还将通过一个多生产者单消费者通道连接这两个任务。在 Rust以及其他语言中同步通道基本上可以实现线程安全和异步的线程间通信从而使其能够更好地扩展。主要流程很简单每当一个嵌入任务完成将文本文档嵌入向量后它会将向量及其 ID文件名“发送”到通道并立即继续处理下一个文档见下图中的读取端。与此同时写入任务会不断地从通道中读取数据将向量分块存储在内存中并在达到一定大小时刷新数据。因为我预计嵌入任务会更加耗时和资源所以我们会将其并行化利用运行应用程序的机器上可用的多个核心。换句话说我们将有多个嵌入任务来读取和嵌入文档以及一个单独的写入任务来分块并将向量写入数据库。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/50e34f48e49edcc6c561c287312198d8.png流水线设计与应用流程图片由作者提供让我们从main()函数开始这将使得管道的流程更加清晰。如上所示在设置好通道第 3 行后我们初始化了写入任务线程该线程开始从通道中轮询消息直到通道关闭。接着它列出了相关目录中的文件并将它们存储在一个字符串集合中。最后它使用Rayon通过par_iter函数并行处理文件列表以便使用process_text_file()函数对其进行并行化处理。使用 Rayon 将允许我们根据机器的能力尽可能扩展文档的并行处理。如你所见流程相对简单主要协调两个主要任务文档处理和向量存储。这个设计允许高效的并行化和可扩展性。文档处理任务使用 Rayon 来并行化文件处理最大化利用可用的系统资源。同时存储任务负责高效地将嵌入的向量写入 LanceDB。关注点的分离不仅简化了整体架构还允许对每个任务进行独立优化。在接下来的章节中我们将更详细地探讨这两个函数。4. 使用 Candle 进行文档嵌入正如我们之前所看到的在我们的管道的一端我们有多个嵌入任务每个任务都在自己的线程上运行。Rayon 的iter_par函数有效地遍历文件列表在每个文件上调用process_text_file()函数同时最大化并行化。让我们从函数本身开始该函数首先获取对嵌入模型的引用这是函数中最棘手的部分我稍后会详细讲解。接下来它将文件分成一定大小的块并对每个块调用嵌入函数该函数实际上调用的是模型本身。嵌入函数返回一个类型为Vecf32的向量大小为[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()的调用将为我们提供模型的线程本地引用如果模型尚未初始化则初始化代码将首先执行。现在我们已经了解了如何运行多个并行执行的嵌入任务并将向量写入通道我们可以继续处理管道的另一部分——写入任务。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/cd944fb01c76ea59b2cda960c7e703a1.png图片来源SpaceX 通过 Unsplash4. 写入任务高效的向量存储写入任务相对简单主要作为一个接口封装了 LanceDB 的写入功能。回想一下LanceDB 是一个嵌入式数据库这意味着它作为一个库的查询引擎读取和写入可以存储在远程存储上的数据例如 AWS S3并且它不拥有数据。这使得它在需要处理大规模数据且低延迟的用例中尤其方便而无需管理单独的数据库服务器。LanceDB 的 Rust API 使用 Arrow 来定义 schema 和表示数据其 Python API 对某些人来说可能更方便。例如以下是我们如何在 Arrow 格式中定义 schema如你所见我们当前的 schema 包含两个字段“filename”字段它将保存实际的文件位置并作为我们的键以及“vector”字段它保存实际的文档向量。在 LanceDB 中向量使用FixedSizeListArrow 类型表示表示一个数组而向量中的每个项目将是 Float32 类型向量的长度最后设置将是 384。连接到 LanceDB 非常简单只需要一个存储位置可以是本地存储路径或 S3 URI。然而使用 Rust 和 Arrow 数据结构将数据附加到 LanceDB 上并不十分开发者友好。与其他基于 Arrow 的列式数据结构类似插入数据时不是附加一行行的数据列表而是每一列都表示为一个值的列表。例如如果你有 10 行数据需要插入且有 2 列你需要附加 2 个列表每个列表包含 10 个值。下面是一个例子代码的核心在第 2 行我们从我们的 schema 和列数据构建一个 Arrow 的RecordBatch。在这个例子中我们有两列——文件名和向量。我们使用两个列表初始化我们的记录批次key_array一个包含文件名的字符串列表以及vectors_array一个包含向量数组的列表。从这里开始Rust 的严格类型安全要求我们在将数据传递给第 1 行获得的表引用的add()函数之前必须对数据进行大量的包装。为了简化这一逻辑我们创建了一个存储模块封装了这些操作并提供一个基于connect(uri)函数和add_vector函数的简单接口。下面是写入任务线程的完整代码该线程从通道读取嵌入数据分块并在达到一定大小时进行写入一旦数据写入LanceDB 数据文件可以从任何进程中访问。下面是一个例子展示如何使用相同的数据进行向量相似性搜索使用的是 LanceDB 的 Python API且该 API 可以从完全不同的进程中执行。uridata/vecdb1dblancedb.connect(uri)tbldb.open_table(vectors_table_1)# the vector we are finding similarities forencoded_vecget_some vector()# perform a similiarity search for top 3 vectorstbl.search(embeddings[0])\.select([filename])\.limit(3).to_pandas()https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/a82402bd8f2bfea7c2a6a5e986dedef9.png脚本输出图片由作者提供5. 结论在这篇文章中我们看到了一个使用 Rust、HuggingFace 的 Candle 框架和 LanceDB 的高性能 RAG 管道的工作示例。我们看到了如何将 Rust 的性能优势与 Candle 相结合来高效地并行读取和嵌入多个文本文件。我们还看到了如何利用同步通道同时运行嵌入任务并与写入流程协同工作而无需处理复杂的锁和同步机制。最后我们学会了如何使用 Rust 利用 LanceDB 的高效存储生成可以与多个 AI 框架和查询库集成的向量存储。我相信这里概述的方法可以作为构建可扩展、生产就绪的 RAG 索引管道的强大基础。无论你是在处理大量数据需要频繁更新知识库还是在资源受限的环境中操作本文讨论的构建块和设计原则都可以根据你的具体需求进行调整。随着 AI 领域的不断发展高效地处理和检索相关信息的能力将始终至关重要。通过结合合适的工具和周到的设计正如本文所展示的开发人员可以创建不仅满足当前需求而且能够应对未来 AI 驱动的信息检索和生成挑战的 RAG 管道。备注与链接GitHub 上的源代码可以在这里找到。该仓库还包含了一个示例 jupyter notebook展示了如何使用 Python 测试这一方法。我之前关于 HuggingFace Candle 的文章可以在这里找到。Candle 框架及其文档包括他们的完整示例文件夹LanceDB及其Rust API 文档