获取批量文本嵌入预测结果

批量获取响应是高效发送大量非延迟时间敏感类嵌入请求的方法。与获取在线响应（一次只能有一个输入请求）不同，您可以在单个批量请求中发送大量 LLM 请求。与对 Vertex AI 中的表格数据执行批量预测的方式类似，您可以确定输出位置、添加输入，然后响应会异步填充到输出位置。

支持批量预测的文本嵌入模型

文本嵌入模型的所有稳定版均支持批量预测，但 textembedding-gecko-multilingual@001 除外。稳定版是指不再处于预览版阶段且在生产环境中完全受支持的版本。如需查看支持的嵌入模型的完整列表，请参阅嵌入模型和版本。

准备输入

批量请求的输入是可以存储在 BigQuery 表中或作为 JSON 行 (JSONL) 文件存储在 Cloud Storage 中的提示列表。每个请求最多可包含 30,000 项提示。

JSONL 示例

本部分介绍如何设置 JSONL 输入和输出的格式。

JSONL 输入示例

{"content":"Give a short description of a machine learning model:"}
{"content":"Best recipe for banana bread:"}

JSONL 输出示例

{"instance":{"content":"Give..."},"predictions": [{"embeddings":{"statistics":{"token_count":8,"truncated":false},"values":[0.2,....]}}],"status":""}
{"instance":{"content":"Best..."},"predictions": [{"embeddings":{"statistics":{"token_count":3,"truncated":false},"values":[0.1,....]}}],"status":""}

BigQuery 示例

本部分介绍如何设置 BigQuery 输入和输出的格式。

BigQuery 输入示例

此示例展示了一个单列 BigQuery 表。

内容
“提供机器学习模型的简短说明：”
“最适合香蕉面包的食谱：”

BigQuery 输出示例

内容	预测	status
“提供机器学习模型的简短说明：”	'[{"embeddings": { "statistics":{"token_count":8,"truncated":false}, "Values":[0.1,....] } } ]'
“最适合香蕉面包的食谱：”	'[{"embeddings": { "statistics":{"token_count":3,"truncated":false}, "Values":[0.2,....] } } ]'

请求批量响应

批量生成任务可能需要一些时间才能完成，具体取决于提交的输入数据项数量。

REST

如需使用 Vertex AI API 测试文本提示，请向发布方模型端点发送 POST 请求。

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

PROJECT_ID：您的 Google Cloud 项目的 ID。
BP_JOB_NAME：作业名称。
INPUT_URI：输入源 URI。这是 BigQuery 表 URI 或 Cloud Storage 中的 JSONL 文件 URI。
OUTPUT_URI：输出目标 URI。

HTTP 方法和网址：

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

请求 JSON 正文：

{
    "name": "BP_JOB_NAME",
    "displayName": "BP_JOB_NAME",
    "model": "publishers/google/models/textembedding-gecko",
    "inputConfig": {
      "instancesFormat":"bigquery",
      "bigquerySource":{
        "inputUri" : "INPUT_URI"
      }
    },
    "outputConfig": {
      "predictionsFormat":"bigquery",
      "bigqueryDestination":{
        "outputUri": "OUTPUT_URI"
    }
  }
}

如需发送请求，请选择以下方式之一：

curl

注意：以下命令假定您已使用您的用户账号通过运行 gcloud init 或 gcloud auth login 登录 gcloud CLI，或者使用了 Cloud Shell，这会使您自动登录 gcloud CLI。您可以运行 gcloud auth list 来检查当前活跃的账号。

将请求正文保存在名为 request.json 的文件中，然后执行以下命令：

curl -X POST \
     -H "Authorization: Bearer $(gcloud auth print-access-token)" \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/batchPredictionJobs"

PowerShell

注意：以下命令假定您已使用您的用户账号通过运行 gcloud init 或 gcloud auth login 登录 gcloud CLI。您可以运行 gcloud auth list 来检查当前活跃的账号。

将请求正文保存在名为 request.json 的文件中，然后执行以下命令：

$cred = gcloud auth print-access-token
$headers = @{ "Authorization" = "Bearer $cred" }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/batchPredictionJobs" | Select-Object -Expand Content

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

{
  "name": "projects/123456789012/locations/us-central1/batchPredictionJobs/1234567890123456789",
  "displayName": "BP_sample_publisher_BQ_20230712_134650",
  "model": "projects/{PROJECT_ID}/locations/us-central1/models/textembedding-gecko",
  "inputConfig": {
    "instancesFormat": "bigquery",
    "bigquerySource": {
      "inputUri": "bq://project_name.dataset_name.text_input"
    }
  },
  "modelParameters": {},
  "outputConfig": {
    "predictionsFormat": "bigquery",
    "bigqueryDestination": {
      "outputUri": "bq://project_name.llm_dataset.embedding_out_BP_sample_publisher_BQ_20230712_134650"
    }
  },
  "state": "JOB_STATE_PENDING",
  "createTime": "2023-07-12T20:46:52.148717Z",
  "updateTime": "2023-07-12T20:46:52.148717Z",
  "labels": {
    "owner": "sample_owner",
    "product": "llm"
  },
  "modelVersionId": "1",
  "modelMonitoringStatus": {}
}

响应包含批量作业的唯一标识符。您可以使用 BATCH_JOB_ID 轮询批量作业的状态，直到作业 state 为 JOB_STATE_SUCCEEDED。例如：

curl \
  -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/batchPredictionJobs/BATCH_JOB_ID

Gen AI SDK for Python

了解如何安装或更新 Google Gen AI SDK for Python。
如需了解详情，请参阅 Gen AI SDK for Python API 参考文档或 python-genai GitHub 代码库。
设置环境变量以将 Gen AI SDK 与 Vertex AI 搭配使用：

# Replace the `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` values
# with appropriate values for your project.
export GOOGLE_CLOUD_PROJECT=GOOGLE_CLOUD_PROJECT
export GOOGLE_CLOUD_LOCATION=us-central1
export GOOGLE_GENAI_USE_VERTEXAI=True

import time

from google import genai
from google.genai.types import CreateBatchJobConfig, JobState, HttpOptions

client = genai.Client(http_options=HttpOptions(api_version="v1"))
# TODO(developer): Update and un-comment below line
# output_uri = "gs://your-bucket/your-prefix"

# See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.batches.Batches.create
job = client.batches.create(
    model="text-embedding-005",
    # Source link: https://storage.cloud.google.com/cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl
    src="gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl",
    config=CreateBatchJobConfig(dest=output_uri),
)
print(f"Job name: {job.name}")
print(f"Job state: {job.state}")
# Example response:
# Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
# Job state: JOB_STATE_PENDING

# See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob
completed_states = {
    JobState.JOB_STATE_SUCCEEDED,
    JobState.JOB_STATE_FAILED,
    JobState.JOB_STATE_CANCELLED,
    JobState.JOB_STATE_PAUSED,
}

while job.state not in completed_states:
    time.sleep(30)
    job = client.batches.get(name=job.name)
    print(f"Job state: {job.state}")
    if job.state == JobState.JOB_STATE_FAILED:
        print(f"Error: {job.error}")
        break

# Example response:
# Job state: JOB_STATE_PENDING
# Job state: JOB_STATE_RUNNING
# Job state: JOB_STATE_RUNNING
# ...
# Job state: JOB_STATE_SUCCEEDED

Python 版 Vertex AI SDK

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

import vertexai

from vertexai.preview import language_models

# TODO(developer): Update & uncomment line below
# PROJECT_ID = "your-project-id"
vertexai.init(project=PROJECT_ID, location="us-central1")
input_uri = (
    "gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl"
)
# Format: `"gs://your-bucket-unique-name/directory/` or `bq://project_name.llm_dataset`
output_uri = OUTPUT_URI

textembedding_model = language_models.TextEmbeddingModel.from_pretrained(
    "textembedding-gecko@003"
)

batch_prediction_job = textembedding_model.batch_predict(
    dataset=[input_uri],
    destination_uri_prefix=output_uri,
)
print(batch_prediction_job.display_name)
print(batch_prediction_job.resource_name)
print(batch_prediction_job.state)
# Example response:
# BatchPredictionJob 2024-09-10 15:47:51.336391
# projects/1234567890/locations/us-central1/batchPredictionJobs/123456789012345
# JobState.JOB_STATE_SUCCEEDED

检索批量输出

批量预测任务完成后，输出存储在您在请求中指定的 Cloud Storage 存储桶或 BigQuery 表中。

后续步骤

了解如何获取文本嵌入。