在 Vertex AI 上部署使用 TEI DLC 的嵌入模型

BGE（代表 BAAI 通用嵌入）是 BAAI 发布的一系列嵌入模型的集合，它是一个用于通用嵌入任务的英文基础模型，在 MTEB 排行榜上名列前茅。文本嵌入推理 (TEI) 是 Hugging Face 开发的一个工具包，用于部署和服务开源文本嵌入和序列分类模型；它能够为最流行的模型（包括 FlagEmbedding、Ember、GTE 和 E5）提供高性能的提取。此外，Google Vertex AI 是一个机器学习 (ML) 平台，它允许您训练和部署 ML 模型和 AI 应用程序，并自定义大型语言模型 (LLM) 以用于您的 AI 驱动的应用程序。

此示例展示了如何在 Vertex AI 上部署任何受支持的嵌入模型（在本例中为 BAAI/bge-large-en-v1.5），使用 Google Cloud Platform (GCP) 中可用的 TEI DLC 在 CPU 和 GPU 实例中从 Hugging Face Hub 部署。

设置/配置

首先，您需要在本地机器上安装 gcloud，它是 Google Cloud 的命令行工具，请按照 Cloud SDK 文档 - 安装 gcloud CLI 中的说明进行操作。

然后，您还需要安装 google-cloud-aiplatform Python SDK，它用于以编程方式创建 Vertex AI 模型、注册模型、创建端点并在 Vertex AI 上部署模型。

!pip install --upgrade --quiet google-cloud-aiplatform

可选地，为了方便在本教程中使用命令，您需要为 GCP 设置以下环境变量

%env PROJECT_ID=your-project-id
%env LOCATION=your-location
%env CONTAINER_URI=us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cu122.1-4.ubuntu2204

然后，您需要登录到您的 GCP 帐户并将项目 ID 设置为您要在 Vertex AI 上注册和部署模型的项目 ID。

!gcloud auth login
!gcloud auth application-default login  # For local development
!gcloud config set project $PROJECT_ID

登录后，您需要在 GCP 中启用必要的服务 API，例如 Vertex AI API、Compute Engine API 和 Google Container Registry 相关的 API。

!gcloud services enable aiplatform.googleapis.com
!gcloud services enable compute.googleapis.com
!gcloud services enable container.googleapis.com
!gcloud services enable containerregistry.googleapis.com
!gcloud services enable containerfilesystem.googleapis.com

在 Vertex AI 上注册模型

一切设置完成后，您可以通过以下方式使用 google-cloud-aiplatform Python SDK 初始化 Vertex AI 会话

import os
from google.cloud import aiplatform

aiplatform.init(
    project=os.getenv("PROJECT_ID"),
    location=os.getenv("LOCATION"),
)

然后，您可以“上传”模型，即在 Vertex AI 上注册模型。它本身并不是真正的上传，因为模型会在启动时通过 MODEL_ID 环境变量从 Hugging Face Hub 中的 TEI 的 Hugging Face DLC 自动下载，因此上传的只是配置，而不是模型权重。

在进入代码之前，让我们快速回顾一下提供给 upload 方法的参数

• display_name 是将在 Vertex AI 模型注册表中显示的名称。

• serving_container_image_uri 是将用于服务模型的 TEI 的 Hugging Face DLC 的位置。

• serving_container_environment_variables 是在容器运行时将使用的环境变量，因此它们与 text-embeddings-inference 定义的环境变量一致，这些环境变量类似于 text-embeddings-router 参数。此外，TEI 的 Hugging Face DLC 还捕获来自 Vertex AI 的 AIP_ 环境变量，如 Vertex AI 文档 - 预测的自定义容器要求 中所述。

  • MODEL_ID 是 Hugging Face Hub 中模型的标识符，要查看所有 TEI 支持的模型，请查看 https://huggingface.co/models?other=text-embeddings-inference&sort=trending。

• （可选）serving_container_ports 是 Vertex AI 端点将公开的端口，默认为 8080。

有关支持的 aiplatform.Model.upload 参数的更多信息，请查看其 Python 参考：https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_upload。

model = aiplatform.Model.upload(
    display_name="BAAI--bge-large-en-v1.5",
    serving_container_image_uri=os.getenv("CONTAINER_URI"),
    serving_container_environment_variables={
        "MODEL_ID": "BAAI/bge-large-en-v1.5",
    },
    serving_container_ports=[8080],
)
model.wait()

在 Vertex AI 上部署模型

在 Vertex AI 上注册模型后，您需要定义要将模型部署到的端点，然后将模型部署链接到该端点资源。

为此，您需要调用 aiplatform.Endpoint.create 方法来创建一个新的 Vertex AI 端点资源（尚未链接到模型或任何可用的内容）。

endpoint = aiplatform.Endpoint.create(display_name="BAAI--bge-large-en-v1.5-endpoint")

现在，您可以在 Vertex AI 上的端点上部署已注册的模型。

deploy 方法会将先前创建的端点资源与包含服务容器配置的模型链接起来，然后在指定的实例上在 Vertex AI 上部署模型。

在进入代码之前，让我们快速回顾一下提供给 deploy 方法的参数

• endpoint 是模型部署到的端点，它是可选的，默认情况下将设置为模型显示名称，并在其后添加 _endpoint 后缀。
• machine_type、accelerator_type 和 accelerator_count 是定义使用哪个实例的参数，此外，分别定义要使用的加速器和加速器的数量。machine_type 和 accelerator_type 是相互关联的，因此您需要选择一个支持您正在使用的加速器的实例，反之亦然。有关不同实例的更多信息，请访问 Compute Engine 文档 - GPU 机器类型，有关 accelerator_type 命名的更多信息，请访问 Vertex AI 文档 - MachineSpec。

有关支持的 aiplatform.Model.deploy 参数的更多信息，您可以在其 Python 参考中查看 https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_deploy。

deployed_model = model.deploy(
    endpoint=endpoint,
    machine_type="g2-standard-4",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
    sync=True,
)

在 Vertex AI 上进行在线预测

最后，您可以使用 predict 方法在 Vertex AI 上运行在线预测，该方法会将请求发送到容器中指定的 /predict 路由中的正在运行的端点，遵循 Vertex AI I/O 的有效负载格式。

output = deployed_model.predict(instances=[{"inputs": "What is Deep Learning?"}])
print(output.predictions[0][0][:5], len(output.predictions[0][0])

这会产生以下输出（为简洁起见已截断，但原始张量的长度为 1024，这是 BAAI/bge-large-en-v1.5 的嵌入维度）。

([0.018108826, 0.0029993146, -0.04984767, -0.035081815, 0.014210845], 1024)

资源清理

最后，您可以在同一个 Python 会话中以编程方式释放资源，如下所示

• deployed_model.undeploy_all 用于从所有端点取消部署模型。
• deployed_model.delete 用于在 undeploy_all 后优雅地删除模型部署到的端点。
• model.delete 用于从注册表中删除模型。

deployed_model.undeploy_all()
deployed_model.delete()
model.delete()