docker-compose up -d

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()
);

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)

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)

create_logger() -> logging.Logger

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

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

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

save_embeddings_to_vectordb(chunks, embeddings) -> None:

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

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

MODELS_BUCKET = 'hf-models'
EMBEDDING_MODEL = 'intfloat/multilingual-e5-small' # 用于将文本块转换为向量嵌入的嵌入模型。
EMBEDDING_MODEL_REVISION = 'ffdcc22a9a5c973ef0470385cef91e1ecb461d9f'

eu.upload_model_to_minio(MODELS_BUCKET, EMBEDDING_MODEL, EMBEDDING_MODEL_REVISION)

CHUNK_SIZE = 1000         # 将被转换为向量嵌入的文本块大小
CHUNK_OVERLAP = 10
DIMENSION = 384           # 嵌入大小

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)

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)
SentenceTransformer 类使使用嵌入模型变得容易。此外，我们正在使用 RecursiveCharacterTextSplitter 来自 LangChain 的工具来拆分或分块我们的文档。它递归地根据字符列表（我们使用其默认列表）拆分文本，从列表中的第一个字符开始，如果第一个拆分太大，则继续下一个字符。目标是将相关的文本片段保持在一起，保留它们的语义关系。所有这些设置工作仅在创建 Embed 对象时执行一次。我们本可以使用一个简单的函数来分发工作，但这种设置工作必须针对每个批次完成，而这在您需要进行设置工作时不是正确的设计。

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

在我们的演示中，我们正在创建一个本地 Ray 实例。我没有使用 Kubernetes 集群。这是在迁移到真实集群之前使代码工作的最佳方法。我们还没有创建任何 Ray Actor，但我们正在发送 Ray 配置信息，这些信息告诉 Ray 每个 Actor 需要哪些环境变量和库。

接下来，我们创建一个 Ray 数据集来保存我们想要发送到 Embed 类实例的所有数据，这些实例将在每个 Ray Actor 中运行。在我们的例子中，每个 Ray Actor 将接收存储在 MinIO 中的对象引用列表（文档的路径）。我们将使用我们上面描述的函数库中的 “get_object_list” 函数。从 “ray.data.from_items()” 返回的 Ray 数据集包含逻辑，该逻辑将把此列表转换为较小的批次，以便在启动分布式嵌入过程时将其发送到每个 Actor。

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

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

最后，我们准备启动分布式嵌入作业。您可能已经注意到，之前的命令运行得非常快。这是因为还没有发生任何计算。Ray 中的转换 (map_batch 被认为是一种转换) 是 “延迟的”。它们只有在您通过迭代数据集、保存数据集或检查数据集的属性来触发数据消耗时才会执行。因此，我们需要向 ds_embed 请求我们 Actor 的返回值。这是在下面完成的。下面的代码段将需要一些时间来运行。

就是这样。我们完成了。一旦上面的代码完成，您将看到类似于下面显示的输出。

总结

在这篇文章中，我们构建了一个分布式嵌入子系统，它可以在工程工作站和完全分布式的云原生生产环境中运行。所呈现的代码具有以下直接解决我们在介绍中确定的复杂性和现实世界问题的优势。

• 可以有效地运行实验，从而可以彻底测试不同的配置选项。
• 除了配置选项之外，还应尝试解析选项。这将使您能够处理多种文件类型并处理文档中的非文本数据。
• 当使用此处显示的代码时，您的生产环境将运行工程师用于测试和实验的相同代码。
• 分布式嵌入子系统可以在集群中运行。集群可以快速扩展以处理需要作为批处理处理的大量文档，并且还可以针对实时工作负载进行扩展。
• 本文中介绍的代码封装了向量数据库调用，使工程师可以交换不同的产品。
• MinIO 是生成式 AI 的最佳存储解决方案。正如我们在本文中看到的那样，嵌入模型和文档必须存储在 高速、可扩展的存储解决方案中。

如果您有任何问题，请务必在 Slack 上联系我们！