对文本嵌入调优

本页面介绍了如何对文本嵌入模型(如 textembedding-geckotextembedding-gecko-multilingual)进行调优。

基础嵌入模型是基于大量文本数据集进行预训练的,可为许多任务提供强大的基准。对于需要专业知识或高度定制性能的场景,模型调优可让您使用自己的相关数据来微调模型的表示法。textembedding-geckotextembedding-gecko-multilingual 模型的稳定版本支持调优。

文本嵌入模型支持监督式调优。 监督式调优使用有标签样本,这些样本展示了您希望在推理期间从文本嵌入模型获得的输出类型。

如需详细了解模型调优,请参阅模型调优的工作原理

预期的质量改进

Vertex AI 使用参数高效调优方法进行自定义。在对公共检索基准数据集执行的实验中,此方法表现出高达 41%(平均 12%)的显著质量提升。

对嵌入模型调优的应用场景

通过对文本嵌入模型调优,您的模型可以适应特定领域或任务的嵌入。如果预训练的嵌入模型不能很好地满足您的特定需求,这将非常有用。例如,您可以针对公司的特定客户服务工单数据集微调嵌入模型。这有助于聊天机器人了解客户通常提出的各种客户服务问题,从而能够更有效地回答问题。如果不进行调优,模型将无法了解客户服务工单的具体细节,也不知道如何解决产品的具体问题。

调优工作流

Vertex AI 上 textembedding-geckotextembedding-gecko-multilingual 的模型调优工作流如下:

  • 准备模型调优数据集。
  • 将模型调优数据集上传到 Cloud Storage 存储桶。
  • 为 Vertex AI Pipelines 配置项目。
  • 创建模型调优作业。
  • 将调优后的模型部署到同名的 Vertex AI 端点。与文本或 Codey 模型调优作业不同,文本嵌入调优作业不会将调优后的模型部署到 Vertex AI 端点。

准备嵌入数据集

用于对嵌入模型调优的数据集包括与您希望模型执行的任务相配的数据。

用于对嵌入模型调优的数据集格式

训练数据集包含以下文件,这些文件需要位于 Cloud Storage 中。文件的路径在启动调优流水线时由参数定义。这三种类型的文件是语料库文件查询文件标签。只有训练标签是必需的,但您也可以提供验证和测试标签,以实现更好的控制。

  • 语料库文件:路径由参数 corpus_path 定义。这是一个 JSONL 文件,其中每一行都包含具有字符串值的字段 _idtitletext_idtext 是必需的,而 title 是可选的。下面是一个示例 corpus.jsonl 文件:

    {"_id": "doc1", "title": "Get an introduction to generative AI on Vertex AI", "text": "Vertex AI Studio offers a Google Cloud console tool for rapidly prototyping and testing generative AI models. Learn how you can use Vertex AI Studio to test models using prompt samples, design and save prompts, tune a foundation model, and convert between speech and text."}
    {"_id": "doc2", "title": "Use gen AI for summarization, classification, and extraction", "text": "Learn how to create text prompts for handling any number of tasks with Vertex AI's generative AI support. Some of the most common tasks are classification, summarization, and extraction. Vertex AI's PaLM API for text lets you design prompts with flexibility in terms of their structure and format."}
    {"_id": "doc3", "title": "Custom ML training overview and documentation", "text": "Get an overview of the custom training workflow in Vertex AI, the benefits of custom training, and the various training options that are available. This page also details every step involved in the ML training workflow from preparing data to predictions."}
    {"_id": "doc4", "text": "Text embeddings are useful for clustering, information retrieval, retrieval-augmented generation (RAG), and more."}
    {"_id": "doc5", "title": "Text embedding tuning", "text": "Google's text embedding models can be tuned on Vertex AI."}
    
  • 查询文件:查询文件包含您的示例查询。该路径由参数 queries_path 定义。查询文件采用 JSONL 格式,并且与语料库文件具有相同的字段。下面是一个示例 queries.jsonl 文件:

    {"_id": "query1", "text": "Does Vertex support generative AI?"}
    {"_id": "query2", "text": "What can I do with Vertex GenAI offerings?"}
    {"_id": "query3", "text": "How do I train my models using Vertex?"}
    {"_id": "query4", "text": "What is a text embedding?"}
    {"_id": "query5", "text": "Can text embedding models be tuned on Vertex?"}
    {"_id": "query6", "text": "embeddings"}
    {"_id": "query7", "text": "embeddings for rag"}
    {"_id": "query8", "text": "custom model training"}
    {"_id": "query9", "text": "Google Cloud PaLM API"}
    
  • 训练标签:路径由参数 train_label_path 定义。train_label_path 是训练标签数据位置的 Cloud Storage URI,在创建调优作业时指定。标签必须是带有标题的 TSV 文件。您的训练标签文件中必须包含一部分查询和语料库。该文件必须包含 query-idcorpus-idscore 列。query-id 是一个字符串,与查询文件中的 _id 键匹配,corpus-id 是一个字符串,与语料库文件中的 _id 匹配。Score 是非负整数值。 如果一对查询和文档不相关,您可以将其从训练标签文件中排除,也可以将其包含在内,并将其得分设为零。任何大于零的分数都表示文档与查询相关。数字越大表示相关性越高。如果省略得分,则默认值为 1。下面是一个示例 train_labels.tsv 文件:

    query-id  corpus-id   score
    query1    doc1    1
    query2    doc2    1
    query3    doc3    2
    query3    doc5  1
    query4    doc4  1
    query4    doc5  1
    query5    doc5  2
    query6    doc4  1
    query6    doc5  1
    query7    doc4  1
    query8    doc3  1
    query9    doc2  1
    
  • 测试标签:可选。测试标签的格式与训练标签的格式相同,并由 test_label_path 参数指定。如果未提供 test_label_path,则测试标签将自动与训练标签拆分。

  • 验证标签:可选。验证标签的格式与训练标签的格式相同,并由 validation_label_path 参数指定。如果未提供 validation_label_path,则验证标签将自动与训练标签拆分。

数据集大小要求

提供的数据集文件必须满足以下限制条件:

  • 查询数必须介于 9 到 10,000 之间。
  • 语料库中的文档数必须介于 9 到 500,000 之间。
  • 每个数据集标签文件必须至少包含 3 个查询 ID,并且在所有数据集拆分中必须至少有 9 个查询 ID。
  • 标签总数必须小于 500,000。

为 Vertex AI Pipelines 配置项目

使用 Vertex AI Pipelines 平台在项目中执行调优。

配置权限

流水线在两个服务代理下执行训练代码。这些服务代理必须获得特定角色,才能开始使用您的项目和数据集进行训练。

Compute Engine 默认服务账号
PROJECT_NUMBER-compute@developer.gserviceaccount.com

此服务账号需要:

  • 对您在 Cloud Storage 中创建的每个数据集文件的 Storage Object Viewer 权限。

  • 对流水线的输出 Cloud Storage 目录 PIPELINE_OUTPUT_DIRECTORYStorage Object User 权限。

  • 对项目的 Vertex AI User 权限。

您可以不使用 Compute Engine 默认服务账号,而是指定自定义服务账号。如需了解详情,请参阅配置具有精细权限的服务账号

Vertex AI Tuning Service Agent
service-PROJECT_NUMBER@gcp-sa-aiplatform-ft.iam.gserviceaccount.com

此服务账号需要:

  • 对您在 Cloud Storage 中创建的每个数据集文件的 Storage Object Viewer 权限。

  • 对流水线的输出 Cloud Storage 目录 PIPELINE_OUTPUT_DIRECTORYStorage Object User 权限。

如需详细了解如何配置 Cloud Storage 数据集权限,请参阅为流水线工件配置 Cloud Storage 存储桶

使用加速器

调优需要 GPU 加速器。以下任一加速器都可用于文本嵌入调优流水线:

  • NVIDIA_L4
  • NVIDIA_TESLA_A100
  • NVIDIA_TESLA_T4
  • NVIDIA_TESLA_V100
  • NVIDIA_TESLA_P100

如要启动调优作业,您选择的加速器类型和区域需要有足够的 Restricted image training GPUs 配额,例如 Restricted image training Nvidia V100 GPUs per region。如需提高项目的配额,请参阅申请更多配额

并非每个区域都提供所有加速器。如需了解详情,请参阅在 Vertex AI 中使用加速器

创建嵌入模型调优作业

您可以使用 Google Cloud 控制台、REST API 或客户端库创建嵌入模型调优作业。

REST

如需创建嵌入模型调优作业,请使用 projects.locations.pipelineJobs.create 方法

在使用任何请求数据之前,请先进行以下替换:

  • PROJECT_ID:您的 Google Cloud 项目 ID。
  • PIPELINE_OUTPUT_DIRECTORY:流水线输出工件的路径,以“gs://”开头。

HTTP 方法和网址:

POST https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/pipelineJobs

请求 JSON 正文:

{
  "displayName": "tune_text_embeddings_model_sample",
  "runtimeConfig": {
    "gcsOutputDirectory": "PIPELINE_OUTPUT_DIRECTORY",
    "parameterValues": {
      "corpus_path": "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl",
      "queries_path": "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl",
      "train_label_path": "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv",
      "test_label_path": "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv",
      "base_model_version_id":"text-embedding-004",
      "task_type": "DEFAULT",
      "batch_size": "128",
      "train_steps": "1000",
      "output_dimensionality": "768",
      "learning_rate_multiplier": "1.0"
    }
  },
  "templateUri": "https://us-kfp.pkg.dev/ml-pipeline/llm-text-embedding/tune-text-embedding-model/v1.1.3"
}

如需发送您的请求,请展开以下选项之一:

您应该收到类似以下内容的 JSON 响应:

启动流水线后,通过 Google Cloud 控制台跟踪调节作业的进度。

前往 Google Cloud 控制台

Python

如需了解如何安装或更新 Vertex AI SDK for Python,请参阅安装 Vertex AI SDK for Python。 如需了解详情,请参阅 Python API 参考文档

import re

from google.cloud.aiplatform import initializer as aiplatform_init
from vertexai.language_models import TextEmbeddingModel


def tune_embedding_model(
    api_endpoint: str,
    base_model_name: str = "text-embedding-004",
    corpus_path: str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl",
    queries_path: str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl",
    train_label_path: str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv",
    test_label_path: str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv",
):  # noqa: ANN201
    """Tune an embedding model using the specified parameters.
    Args:
        api_endpoint (str): The API endpoint for the Vertex AI service.
        base_model_name (str): The name of the base model to use for tuning.
        corpus_path (str): GCS URI of the JSONL file containing the corpus data.
        queries_path (str): GCS URI of the JSONL file containing the queries data.
        train_label_path (str): GCS URI of the TSV file containing the training labels.
        test_label_path (str): GCS URI of the TSV file containing the test labels.
    """
    match = re.search(r"^(\w+-\w+)", api_endpoint)
    location = match.group(1) if match else "us-central1"
    base_model = TextEmbeddingModel.from_pretrained(base_model_name)
    tuning_job = base_model.tune_model(
        task_type="DEFAULT",
        corpus_data=corpus_path,
        queries_data=queries_path,
        training_data=train_label_path,
        test_data=test_label_path,
        batch_size=128,  # The batch size to use for training.
        train_steps=1000,  # The number of training steps.
        tuned_model_location=location,
        output_dimensionality=768,  # The dimensionality of the output embeddings.
        learning_rate_multiplier=1.0,  # The multiplier for the learning rate.
    )
    return tuning_job

Java

在尝试此示例之前,请按照《Vertex AI 快速入门:使用客户端库》中的 Java 设置说明执行操作。 如需了解详情,请参阅 Vertex AI Java API 参考文档

如需向 Vertex AI 进行身份验证,请设置应用默认凭据。 如需了解详情,请参阅为本地开发环境设置身份验证

import com.google.cloud.aiplatform.v1.CreatePipelineJobRequest;
import com.google.cloud.aiplatform.v1.LocationName;
import com.google.cloud.aiplatform.v1.PipelineJob;
import com.google.cloud.aiplatform.v1.PipelineJob.RuntimeConfig;
import com.google.cloud.aiplatform.v1.PipelineServiceClient;
import com.google.cloud.aiplatform.v1.PipelineServiceSettings;
import com.google.protobuf.Value;
import java.io.IOException;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class EmbeddingModelTuningSample {
  public static void main(String[] args) throws IOException {
    // TODO(developer): Replace these variables before running this sample.
    String apiEndpoint = "us-central1-aiplatform.googleapis.com:443";
    String project = "PROJECT";
    String baseModelVersionId = "BASE_MODEL_VERSION_ID";
    String taskType = "DEFAULT";
    String pipelineJobDisplayName = "PIPELINE_JOB_DISPLAY_NAME";
    String outputDir = "OUTPUT_DIR";
    String queriesPath = "QUERIES_PATH";
    String corpusPath = "CORPUS_PATH";
    String trainLabelPath = "TRAIN_LABEL_PATH";
    String testLabelPath = "TEST_LABEL_PATH";
    double learningRateMultiplier = 1.0;
    int outputDimensionality = 768;
    int batchSize = 128;
    int trainSteps = 1000;

    createEmbeddingModelTuningPipelineJob(
        apiEndpoint,
        project,
        baseModelVersionId,
        taskType,
        pipelineJobDisplayName,
        outputDir,
        queriesPath,
        corpusPath,
        trainLabelPath,
        testLabelPath,
        learningRateMultiplier,
        outputDimensionality,
        batchSize,
        trainSteps);
  }

  public static PipelineJob createEmbeddingModelTuningPipelineJob(
      String apiEndpoint,
      String project,
      String baseModelVersionId,
      String taskType,
      String pipelineJobDisplayName,
      String outputDir,
      String queriesPath,
      String corpusPath,
      String trainLabelPath,
      String testLabelPath,
      double learningRateMultiplier,
      int outputDimensionality,
      int batchSize,
      int trainSteps)
      throws IOException {
    Matcher matcher = Pattern.compile("^(?<Location>\\w+-\\w+)").matcher(apiEndpoint);
    String location = matcher.matches() ? matcher.group("Location") : "us-central1";
    String templateUri =
        "https://us-kfp.pkg.dev/ml-pipeline/llm-text-embedding/tune-text-embedding-model/v1.1.3";
    PipelineServiceSettings settings =
        PipelineServiceSettings.newBuilder().setEndpoint(apiEndpoint).build();
    try (PipelineServiceClient client = PipelineServiceClient.create(settings)) {
      Map<String, Value> parameterValues =
          Map.of(
              "base_model_version_id", valueOf(baseModelVersionId),
              "task_type", valueOf(taskType),
              "queries_path", valueOf(queriesPath),
              "corpus_path", valueOf(corpusPath),
              "train_label_path", valueOf(trainLabelPath),
              "test_label_path", valueOf(testLabelPath),
              "learning_rate_multiplier", valueOf(learningRateMultiplier),
              "output_dimensionality", valueOf(outputDimensionality),
              "batch_size", valueOf(batchSize),
              "train_steps", valueOf(trainSteps));
      PipelineJob pipelineJob =
          PipelineJob.newBuilder()
              .setTemplateUri(templateUri)
              .setDisplayName(pipelineJobDisplayName)
              .setRuntimeConfig(
                  RuntimeConfig.newBuilder()
                      .setGcsOutputDirectory(outputDir)
                      .putAllParameterValues(parameterValues)
                      .build())
              .build();
      CreatePipelineJobRequest request =
          CreatePipelineJobRequest.newBuilder()
              .setParent(LocationName.of(project, location).toString())
              .setPipelineJob(pipelineJob)
              .build();
      return client.createPipelineJob(request);
    }
  }

  private static Value valueOf(String s) {
    return Value.newBuilder().setStringValue(s).build();
  }

  private static Value valueOf(int n) {
    return Value.newBuilder().setNumberValue(n).build();
  }

  private static Value valueOf(double n) {
    return Value.newBuilder().setNumberValue(n).build();
  }
}

Node.js

在尝试此示例之前,请按照《Vertex AI 快速入门:使用客户端库》中的 Node.js 设置说明执行操作。 如需了解详情,请参阅 Vertex AI Node.js API 参考文档

如需向 Vertex AI 进行身份验证,请设置应用默认凭据。 如需了解详情,请参阅为本地开发环境设置身份验证

async function main(
  apiEndpoint,
  project,
  outputDir,
  pipelineJobDisplayName = 'embedding-customization-pipeline-sample',
  baseModelVersionId = 'text-embedding-004',
  taskType = 'DEFAULT',
  corpusPath = 'gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl',
  queriesPath = 'gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl',
  trainLabelPath = 'gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv',
  testLabelPath = 'gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv',
  outputDimensionality = 768,
  learningRateMultiplier = 1.0,
  batchSize = 128,
  trainSteps = 1000
) {
  const aiplatform = require('@google-cloud/aiplatform');
  const {PipelineServiceClient} = aiplatform.v1;
  const {helpers} = aiplatform; // helps construct protobuf.Value objects.

  const client = new PipelineServiceClient({apiEndpoint});
  const match = apiEndpoint.match(/(?<L>\w+-\w+)/);
  const location = match ? match.groups.L : 'us-central1';
  const parent = `projects/${project}/locations/${location}`;
  const params = {
    base_model_version_id: baseModelVersionId,
    task_type: taskType,
    queries_path: queriesPath,
    corpus_path: corpusPath,
    train_label_path: trainLabelPath,
    test_label_path: testLabelPath,
    batch_size: batchSize,
    train_steps: trainSteps,
    output_dimensionality: outputDimensionality,
    learning_rate_multiplier: learningRateMultiplier,
  };
  const runtimeConfig = {
    gcsOutputDirectory: outputDir,
    parameterValues: Object.fromEntries(
      Object.entries(params).map(([k, v]) => [k, helpers.toValue(v)])
    ),
  };
  const pipelineJob = {
    templateUri:
      'https://us-kfp.pkg.dev/ml-pipeline/llm-text-embedding/tune-text-embedding-model/v1.1.3',
    displayName: pipelineJobDisplayName,
    runtimeConfig,
  };
  async function createTuneJob() {
    const [response] = await client.createPipelineJob({parent, pipelineJob});
    console.log(`job_name: ${response.name}`);
    console.log(`job_state: ${response.state}`);
  }

  await createTuneJob();
}

控制台

如需使用 Google Cloud 控制台调优文本嵌入模型,您可以按照以下步骤启动自定义流水线:

  1. 在 Google Cloud 控制台的 Vertex AI 部分中,进入 Vertex AI Pipelines 页面。

    前往 Vertex AI Pipelines

  2. 点击创建运行作业以打开创建流水线运行作业窗格。
  3. 点击从现有流水线中选择,然后输入以下详细信息:
    1. 选择资源下拉列表中选择“ml-pipeline”。
    2. 仓库下拉列表中选择“llm-text-embedding”。
    3. 流水线或组件下拉列表中选择“tune-text-embedding-model”。
    4. 版本下拉列表中选择标记为“v1.1.3”的版本。
  4. 指定运行名称以唯一标识流水线运行。
  5. 区域下拉列表中,选择要创建流水线运行的区域,该区域将与您创建调优模型的区域相同。
  6. 点击继续。此时会显示运行时配置窗格。
  7. Cloud Storage 位置下,点击浏览以选择用于存储流水线输出工件的 Cloud Storage 存储桶,然后点击选择
  8. 流水线参数下,指定用于调优流水线的参数。三个必需参数 corpus_pathqueries_pathtrain_label_path,其格式如准备嵌入数据集中所述。如需详细了解每个参数,请参阅本部分的“REST”标签页。
  9. 点击提交以创建流水线运行。

其他支持的功能

文本嵌入调优支持 VPC Service Controls,并且可以配置为在虚拟私有云 (VPC) 中运行,方法是在创建 PipelineJob 时传递 network 参数。

如需使用 CMEK(客户管理的加密密钥),请将密钥传递给 parameterValues.encryption_spec_key_name 流水线参数,以及 encryptionSpec.kmsKeyName 参数(在创建 PipelineJob 时)。

使用调优后的模型

在 Model Registry 中查看经过调优的模型

调优作业完成后,调优后的模型不会自动部署到端点。它将作为 Model Registry 中的模型资源提供。您可以使用 Google Cloud 控制台查看当前项目中的模型列表,包括已调优的模型。

如需在 Google Cloud 控制台中查看调优后模型,请进入 Vertex AI Model Registry 页面。

进入 Vertex AI Model Registry

部署模型

对嵌入模型调优后,您需要部署模型资源。如需部署调优后的嵌入模型,请参阅将模型部署到端点

与基础模型不同,调优后的文本嵌入模型由用户管理。这包括管理机器类型和加速器等服务资源。为防止在预测期间出现内存不足错误,建议您使用 NVIDIA_TESLA_A100 GPU 类型进行部署,该类型可支持任何输入长度大小不超过 5 的批量大小。

textembedding-gecko 基础模型类似,调优后的模型支持多达 3072 个词元,并且可以截取更长的输入。

使用部署的模型获取预测结果

部署经过调优的模型后,您可以使用以下命令之一向调优后的模型端点发出请求。

调优后的 textembedding-gecko@001 模型的 curl 命令示例

如需使用 textembedding-gecko@001 的调优后版本获取预测结果,请使用以下示例 curl 命令。

PROJECT_ID=PROJECT_ID
LOCATION=LOCATION
ENDPOINT_URI=https://${LOCATION}-aiplatform.googleapis.com
MODEL_ENDPOINT=TUNED_MODEL_ENDPOINT_ID

curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json"  \
    ${ENDPOINT_URI}/v1/projects/${PROJECT_ID}/locations/${LOCATION}/endpoints/${MODEL_ENDPOINT}:predict \
    -d '{
  "instances": [
    {
      "content": "Dining in New York City"
    },
    {
      "content": "Best resorts on the east coast"
    }
  ]
}'

textembedding-gecko@001 模型的 curl 命令示例

其他模型的调优后版本(例如 textembedding-gecko@003textembedding-gecko-multilingual@001)还需要另外两个输入:task_typetitle。如需详细了解这些参数,请参阅 curl 命令

PROJECT_ID=PROJECT_ID
LOCATION=LOCATION
ENDPOINT_URI=https://${LOCATION}-aiplatform.googleapis.com
MODEL_ENDPOINT=TUNED_MODEL_ENDPOINT_ID

curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json"  \
    ${ENDPOINT_URI}/v1/projects/${PROJECT_ID}/locations/${LOCATION}/endpoints/${MODEL_ENDPOINT}:predict \
    -d '{
  "instances": [
    {
      "content": "Dining in New York City",
      "task_type": "DEFAULT",
      "title": ""
    },
    {
      "content": "There are many resorts to choose from on the East coast...",
      "task_type": "RETRIEVAL_DOCUMENT",
      "title": "East Coast Resorts"
    }
  ]
}'

输出示例

无论什么版本,此输出都适用于 textembedding-geckotextembedding-gecko-multilingual 模型。

{
 "predictions": [
   [ ... ],
   [ ... ],
   ...
 ],
 "deployedModelId": "...",
 "model": "projects/.../locations/.../models/...",
 "modelDisplayName": "tuned-text-embedding-model",
 "modelVersionId": "1"
}

后续步骤