使用 MinIO、Langchain 和 Ray Data 构建分布式嵌入式子系统

嵌入子系统是实现检索增强生成所需的四个子系统之一。它将您的自定义语料库转换为可以搜索语义含义的向量数据库。其他子系统是用于创建自定义语料库的数据管道，用于查询向量数据库以向用户查询添加更多上下文的检索器，最后是托管大型语言模型（LLM）的服务子系统，并将根据用户的查询和在向量数据库中找到的上下文生成答案。下图显示了这四个子系统在检索增强生成过程中如何协同工作。

在这篇文章中，我想重点介绍嵌入子系统。在此子系统中，构成组织自定义语料库的文档从其原始格式转换为文本，分解为较小的块，然后为每个块创建一个嵌入（这是一个向量，通常具有数百个维度）。创建嵌入后，原始块和向量都将存储在向量数据库中。嵌入子系统在概念上易于理解，并且实现嵌入简单文本文件的简单脚本非常简单。但是，如果您必须为您的组织实施嵌入子系统，那么您如何为您的组织做出正确的设计决策，以及您如何应对不断增长的需求带来的复杂性？下面列出了一些设计决策和实际复杂性：

如何高效地运行多个实验来测试不同的配置选项？
如何处理文档中的表格和图像？
如何将嵌入子系统部署到生产环境中？
如何处理需要嵌入的大量文档？
什么是最好的载体数据库？
文档、嵌入模型和LLMs的最佳存储选项是什么？

解决这些问题的第一步是使用能够在工程工作站以及生产环境中运行的现代工具。具体来说，我们将使用 MinIO 进行所有存储，使用 Langchain 作为低代码解决方案来进行文档解析（我还将提供一些比 Langchain 更好地处理图像和表的选项），并使用 Ray Data 将分块和嵌入函数分发到集群中。毫不奇怪，分布式技术是我们解决方案的基础。您不仅可以使用商用硬件设置进行并行处理来获得高吞吐量，而且该解决方案是云原生的，使其可以跨云移植，并且还能够在本地运行。让我们从为我们的实验设置一个自定义语料库开始。

在 MinIO 中设置自定义语料库

如上所述，自定义语料库由数据管道创建，该数据管道将可能位于组织中多个门户的文档聚合到 MinIO 中。创建文档管道是另一篇文章的主题 - 所以现在，我们将手动将一些文档暂存到 MinIO 桶中。我在这里也只会使用文本文档来保持简单。但是，这里有一些处理文档中多种文件格式和非文本的提示。首先，查看 Unstructured 的库，用于分区、清理和提取。其次，如果您专门处理 PDF，请查看 Open-Parse 库。我们在之前的博客文章《使用 Open-Parse 智能分块提高 RAG 性能》中介绍了 Open-Parse。下面的屏幕截图显示了我们的自定义语料库。我从古腾堡计划中下载了四本被认为是经典的流行书籍的文本版本。

人性论——大卫·休谟
孙子兵法 - 孙子
杰基尔博士和海德先生的奇案 - 罗伯特·路易斯·史蒂文森
《海底两万里》——儒勒·凡尔纳

现在我们有了一个自定义语料库，我们可以设置一个向量数据库来保存嵌入。

设置 MinIO 和矢量数据库

我将使用的向量数据库是 Pgvector。Pgvector 是 PostgreSQL 的开源扩展，允许用户在数据库中存储、搜索和分析矢量数据。这篇文章的代码下载有一个 docker-compose 文件，其中包含 MinIO、Pgvector 和 pgAdmin。在与 docker-compose fill 相同的目录中运行以下命令会将这三个服务作为容器显示出来。

docker-compose up -d

还有一个init.sql文件（如下所示）。docker-compose 文件将此文件映射到容器的启动目录。这会导致文件中的 SQL 运行，从而在 Postgres 中创建向量扩展和一个“嵌入”表，其中包含下面 SQL 文件中显示的字段。

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS embeddings (
 id SERIAL PRIMARY KEY,
 embedding vector,
 text text,
 created_at timestamptz DEFAULT now()
);

将嵌入模型保存到 MinIO

我们将使用的嵌入模型是 Hugging Face 的开源模型。详细信息如下所示。在运行实验时指定特定版本始终是一个好主意。模型名称：intfloat/multilingual-e5-small修订版：ffdcc22a9a5c973ef0470385cef91e1ecb461d9f

不要被模型的名字所迷惑。它一点也不小。它是 1.4GB。我们需要下载这个模型并上传到 MinIO。这是一项一次性设置任务，用于在分布式环境中暂存此模型以进行实验。遗憾的是，我们需要的 Hugging Face 函数（snapshot_download）没有 S3 接口，所以我们会使用 MinIO Python SDK 将模型上传到 MinIO。更复杂的是，Hugging Face 模型不是单个文件。它是下载到指定目录中的文件集合。我们必须将整个目录上传到 MinIO，并使用 MinIO 路径保留文件夹结构。这是使用如下所示的“upload_model_to_minio”函数完成的。

from huggingface_hub import snapshot_download

def upload_model_to_minio(bucket_name: str, full_model_name: str, revision: str) -> None:
   '''
   Download a model from Hugging Face and upload it to MinIO. This function will use
   the current systems temp directory to temporarily save the model.
   '''

   # Create a local directory for the model.
   #home = str(Path.home())
   temp_dir = tempfile.gettempdir()
   base_path = f'{temp_dir}{os.sep}hf-models'
   os.makedirs(base_path, exist_ok=True)

   # Get the user name and the model name.
   tmp = full_model_name.split('/')
   user_name = tmp[0]
   model_name = tmp[1]

   # The snapshot_download will use this pattern for the path name.
   model_path_name=f'models--{user_name}--{model_name}'
   # The full path on the local drive.
   full_model_local_path = base_path + os.sep + model_path_name + os.sep + 'snapshots' + 
                           os.sep + revision
   # The path used by MinIO.
   full_model_object_path = model_path_name + '/snapshots/' + revision

   print(f'Starting download from HF to {full_model_local_path}.')
   snapshot_download(repo_id=full_model_name, revision=revision, cache_dir=base_path)

   print('Uploading to MinIO.')
   upload_local_directory_to_minio(full_model_local_path, bucket_name, 
                                   full_model_object_path)
   shutil.rmtree(full_model_local_path)

运行以下命令将使用此函数将我们的模型上传到名为“hf-models”的存储桶中。

MODELS_BUCKET = 'hf-models'
EMBEDDING_MODEL = 'intfloat/multilingual-e5-small'
EMBEDDING_MODEL_REVISION = 'ffdcc22a9a5c973ef0470385cef91e1ecb461d9f'

upload_model_to_minio(MODELS_BUCKET, EMBEDDING_MODEL, EMBEDDING_MODEL_REVISION)

嵌入函数库

当你使用像 Ray Data 这样的库来分发数据处理时——在本例中是文本的分块和每个块的嵌入生成——你真正要做的就是编排简单的函数调用，这些函数调用在此过程中执行一项任务。下面列出了从 MinIO 存储桶中的文档列表创建嵌入所需的所有函数，以及它们的参数和返回值。如您所见，我们拥有嵌入文档集合所需的一切。

create_logger() -> logging.Logger

创建一个 Python 记录器，用于将调试、信息、错误、警告和关键消息发送到日志记录存储库。

download_model_from_minio(bucket_name: str, full_model_name: str, revision: str) -> str

将模型从 MinIO 下载到当前系统临时目录。一旦它被加载到内存中，它将删除它。

get_document_from_minio(bucket_name: str, object_name: str) -> str:

从 MinIO 下载单个文档，并将其保存到当前系统临时目录中。

get_object_list(bucket_name: str) -> List[str]:

返回指定存储桶中的对象列表。此列表将发送到 Ray Data，后者将其均匀分布在集群中的所有 Ray actor 中。

save_embeddings_to_vectordb(chunks, embeddings) -> None:

将嵌入和文本块一起保存到向量数据库中。

upload_local_directory_to_minio(local_path:str, bucket_name:str , minio_path:str) -> None

将指定本地目录的内容上传到 MinIO，将文件夹结构保留为指定存储桶内的路径。

upload_model_to_minio(bucket_name: str, full_model_name: str, revision: str) -> None:

从Hugging Face下载模型到当前系统临时目录，然后将模型上传到指定的Bucket，同时保留文件夹结构作为指定Bucket内的路径。

一个简单的嵌入子系统

让我们使用上述函数并创建一个简单的非分布式脚本。下面的代码将为 Robert Louis Stevenson 的“The Strange Case of Dr Jekyll and Mr Hyde”创建嵌入。首先，我们需要下载我们希望使用的嵌入模型，并将其保存到 MinIO 中。这是一项一次性任务;您无需在每次想要嵌入新一批模型或运行实验时都这样做。

MODELS_BUCKET = 'hf-models'
EMBEDDING_MODEL = 'intfloat/multilingual-e5-small' # Embedding model to use for converting text chunks to vector embeddings.
EMBEDDING_MODEL_REVISION = 'ffdcc22a9a5c973ef0470385cef91e1ecb461d9f'

eu.upload_model_to_minio(MODELS_BUCKET, EMBEDDING_MODEL, EMBEDDING_MODEL_REVISION)

接下来，我们需要从 MinIO 下载我们的模型，实例化它，创建一个分块器（或拆分器），创建嵌入并将它们保存到我们的 pgvector 数据库中。

CHUNK_SIZE = 1000         # Text chunk sizes which will be converted to vector embeddings
CHUNK_OVERLAP = 10
DIMENSION = 384           # Embeddings size

model_path = eu.download_model_from_minio(MODELS_BUCKET, EMBEDDING_MODEL, 

                                          EMBEDDING_MODEL_REVISION)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
embedding_model = SentenceTransformer(model_path, device=device)
chunker = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, 

                                         chunk_overlap=CHUNK_OVERLAP, length_function=len)

temp_file = eu.get_document_from_minio(BUCKET_NAME, 

                                       'The Strange Case of Dr Jekyll and Mr Hyde.txt')
file = open(temp_file, 'r')
data = file.read()
chunks = chunker.split_text(data)
print('Number of chunks:', len(chunks))
print('Length of the first chunk:', len(chunks[0]))

embeddings = embedding_model.encode(chunks, batch_size=BATCH_SIZE).tolist()
print('Number of embeddings:', len(embeddings))
print('Length of the first embedding:', len(embeddings[0]))

eu.save_embeddings_to_vectordb(chunks, embeddings)

请注意，如果我们可以访问 GPU，则我们正在使用 GPU。此外，一切都是配置驱动的，因此运行不同的实验就是更改配置以反映您希望运行的实验的问题。这包括根据需要更改嵌入模型。

下面是 pgAdmin 的屏幕截图，显示了我们新创建的嵌入。

现在我们有了一个简单的脚本，可以为单个文档创建嵌入，下一步是将此代码迁移到在集群中运行的框架。这将允许并行嵌入整个文档语料库。我们将使用 Ray Data 来执行此操作。

分发嵌入子系统

分发嵌入子系统的第一步是将所有工作放入一个行为类似于函数的类中。这是使用 Python 的“__call__”内置方法完成的。（这是光线数据的要求。我们的班级如下所示。

class Embed:

   def __init__(self):
       device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      
       model_path = eu.download_model_from_minio(MODELS_BUCKET, EMBEDDING_MODEL, 

                                                 EMBEDDING_MODEL_REVISION)
       self.embedding_model = SentenceTransformer(model_path, device=device)
       self.splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, 

                                                      chunk_overlap=CHUNK_OVERLAP, 

                                                      length_function=len)

   def __call__(self, batch_list: List[str]) -> None:
       document_list = batch_list["item"]

       timings = []
       documents = []
       for document_data in document_list:
           start_time = time()
           bucket_name = document_data[0]
           object_name = document_data[1]
           temp_file = eu.get_document_from_minio(bucket_name, object_name)
           file = open(temp_file, 'r')
           data = file.read()

           chunks = self.splitter.split_text(data)
           embeddings = self.embedding_model.encode(chunks, batch_size=BATCH_SIZE).tolist()
           eu.save_embeddings_to_vectordb(chunks, embeddings)
          
           total_time_sec = time() - start_time
           documents.append(object_name)
           timings.append(total_time_sec)

       return {'timings': timings, 'documents': documents}

Ray Data 将为我们将很快创建的 Ray 集群中的每个 actor 实例化一次此类。此对象将保持活动状态，并接收多个批次进行处理。请注意，“__init__”函数正在下载我们的嵌入模型，并使用它创建一个 SentenceTransformer。SentenceTransformer 类使使用嵌入模型变得容易。此外，我们还使用 LangChain 的 RecursiveCharacterTextSplitter 来分割或分块我们的文档。它根据字符列表（我们使用其默认列表）递归拆分文本，从列表中的第一个字符开始，如果第一个拆分太大，则继续到下一个字符。目标是将相关的文本片段保持在一起，保留它们的语义关系。所有这些设置工作仅在创建 Embed 对象时发生一次。我们本可以使用一个简单的函数来分配工作，但是必须为每个批次完成此设置工作，当您要进行设置工作时，这不是正确的设计。

接下来，我们需要初始化 Ray 集群。


ray.init(
   #address="ray://ray-cluster-kuberay-head-svc:10001",
   runtime_env={
       "env_vars": {
           "MINIO_URL": MINIO_URL,
           "MINIO_ACCESS_KEY": MINIO_ACCESS_KEY,
           "MINIO_SECRET_KEY": MINIO_SECRET_KEY,
           "MINIO_SECURE": str(MINIO_SECURE),
           "PGVECTOR_HOST": os.environ['PGVECTOR_HOST'],
           "PGVECTOR_DATABASE": os.environ['PGVECTOR_DATABASE'],
           "PGVECTOR_USER": os.environ['PGVECTOR_USER'],
           "PGVECTOR_PASSWORD": os.environ['PGVECTOR_PASSWORD'],
           "PGVECTOR_PORT": os.environ['PGVECTOR_PORT'],
       },
       "pip": [              
           "datasets==2.19.0",
           "huggingface_hub==0.22.2",
           "minio==7.2.7",
           "psycopg2-binary==2.9.9",
           "pyarrow==16.0.0",
           "sentence-transformers==3.0.1",
           "torch==2.3.0",
           "transformers==4.40.1",
       ]
   }
)

在我们的演示中，我们将创建一个本地 Ray 实例。我没有使用 Kubernetes 集群。这是在移动到真实集群之前让代码正常工作的最佳方法。我们还没有创建任何 Ray actor - 但我们正在发送 Ray 配置信息，告诉 Ray 每个 actor 需要的环境变量和库。接下来，我们创建一个 Ray 数据集来保存我们想要发送到将在每个 Ray actor 中运行的 Embed 类实例的所有数据。在我们的例子中，每个 Ray actor 都将收到一个存储在 MinIO 中的对象引用列表（文档的路径）。我们将使用上述函数库中的“get_object_list”函数。从 “ray.data.from_items（）返回的 Ray 数据集包含的逻辑，当我们启动分布式嵌入过程时，该逻辑会将此列表转换为每个 actor 的较小批次。


# The embedding class expects bucket_name and document_name pairs - so add bucket name to each entry in the list.
document_list = eu.get_object_list(BUCKET_NAME)
list_for_ray = [[BUCKET_NAME, doc] for doc in document_list]

ray_ds = ray.data.from_items(list_for_ray)
print(type(ray_ds))
print(ray_ds.schema)

我们几乎已经准备好进行一些分布式计算，但我们还有一个编码任务要完成。我们需要将 Ray 数据集映射到我们的 Embed 类，并告诉 Ray 如何设置我们之前为此工作负载初始化的集群。这是使用 Ray 数据集的“map_batches”方法完成的。您可以将函数或可调用类发送到“map_batches”。如果发送函数，Ray Data 将使用无状态 Ray 任务。对于类，Ray Data 使用有状态的 Ray actor。

ds_embed = ray_ds.map_batches(
   Embed,
   concurrency=ACTOR_POOL_SIZE,
   batch_size=BATCH_SIZE,  # Large batch size to maximize GPU utilization.
   #num_gpus=1,            # 1 GPU for each actor.
   num_cpus=1,             # 1 CPU for each actor.
)

请注意，我们正在传入需要为每个 actor 实例化的 Embed 类。我们还指定了 actor 的数量、每次调用 actor 的批量大小，最后指定了每个 actor 可访问的 GPU 和 CPU 数量。map_batches 方法返回另一个 Ray 数据集（ds_embed），其中包含所有 actor 的所有返回值。这是 Embedded 中“__call__”方法的返回值的集合。

最后，我们准备开始我们的分布式嵌入工作。您可能已经注意到，上一个命令运行得非常快。那是因为还没有进行任何计算。Ray 中的转换（map_batch被认为是转换）是“惰性”。在通过循环访问数据集、保存数据集或检查数据集的属性来触发数据使用之前，不会执行它们。因此，我们需要向ds_embed请求 actor 的返回值。这是在下面完成的。下面的代码片段需要一些时间才能运行。

def ray_data_task(ds_embed):
   results = []
   for row in ds_embed.iter_rows():
       documents = row['documents']
       timings = row['timings']
       results.append((documents, timings))
   return results

results = ray_data_task(ds_embed)

results

就是这样。大功告成。完成上述代码后，您将看到类似于下面显示的输出。

[('A Treatise of Human Nature.txt', 75.08733916282654),
 ('The Art of War.txt', 21.960258722305298),
 ('The Strange Case of Dr Jekyll and Mr Hyde.txt', 10.052802085876465),
 ('Twenty Thousand Leagues under the Sea.txt', 39.24100613594055)]