获取批量嵌入推理结果

批量获取响应是高效发送大量非延迟时间敏感类嵌入请求的方法。与获取在线响应(一次只能有一个输入请求)不同,您可以在单个批量请求中发送多个嵌入请求。与对 表格数据 在 Gemini Enterprise Agent Platform中执行批量推理的方式类似, 您可以确定输出位置、添加输入,然后响应会 异步填充到输出位置。

Gemini 嵌入批量预测(预览版)

Gemini 嵌入批量推理(预览版)引入了更新的输入和输出架构。每个批量作业最多支持 1,000,000 行,每个项目最多支持队列中的 10,000 个并发批量作业。

如需了解端点和可用位置,请参阅位置

JSONL 输入示例

Gemini 嵌入 001

gemini-embedding-001 模型支持为每一行配置任务类型和标题。

{"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "embed_content_config":{"output_dimensionality": 768, "title": "some_title", "task_type": "RETRIEVAL_DOCUMENT"}}}
{"key": "id_2", "request": {"content": {"parts": [{"text": "Hello World again"}]}, "embed_content_config":{"output_dimensionality": 768, "title": "some_title_2", "task_type": "RETRIEVAL_DOCUMENT"}}}

Gemini 嵌入 2

gemini-embedding-2 模型要求任务类型和标题直接内嵌在提示中。

{"key": "id_1", "request": {"content": {"parts": [{"text": "title: sample | text: Hello World"}]}, "embed_content_config":{"output_dimensionality": 768}}}
{"key": "id_2", "request": {"content": {"parts": [{"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", "mimeType": "image/jpeg"}}]}, "embed_content_config":{"output_dimensionality": 768}}}

JSONL 输出示例

{"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024, ...]}}}
{"key": "id_2", "request": {...}, "response": {"tokenCount": "3", "embedding": {"values": [-0.075, 0.031, ...]}}}

其他资源

如需详细了解如何对 Gemini 嵌入模型使用批量推理,请参阅批量预测 API 参考文档

旧版嵌入批量预测

支持批量推理的文本嵌入模型

所有稳定版旧版文本嵌入模型都支持批量推理,并且在生产环境中完全受支持。如需查看 嵌入模型的完整列表,请参阅嵌入模型和版本

准备输入

批量请求的输入是可以存储在 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

您可以使用 REST 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

将请求正文保存在名为 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

将请求正文保存在名为 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": {}
}
### 示例 curl 命令

请替换命令中的以下变量:

  • PROJECT_ID:您的 Google Cloud 项目 ID。
  • LOCATION:用于发送请求的位置。
  • GCS_SOURCE_BUCKET:您的 Cloud Storage 源存储桶 URI。
  • GCS_OUTPUT_BUCKET:您的 Cloud Storage 输出存储桶 URI。

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

Python

安装

pip install --upgrade google-genai

如需了解详情,请参阅 SDK 参考文档

设置环境变量以将 Google 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_ENTERPRISE=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/.../locations/.../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

Go

了解如何安装或更新 Go

如需了解详情,请参阅 SDK 参考文档

设置环境变量以将 Google 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_ENTERPRISE=True

import (
	"context"
	"fmt"
	"io"
	"time"

	"google.golang.org/genai"
)

// generateBatchEmbeddings shows how to run a batch embeddings prediction job.
func generateBatchEmbeddings(w io.Writer, outputURI string) error {
	// outputURI = "gs://your-bucket/your-prefix"
	ctx := context.Background()

	client, err := genai.NewClient(ctx, &genai.ClientConfig{
		HTTPOptions: genai.HTTPOptions{APIVersion: "v1"},
	})
	if err != nil {
		return fmt.Errorf("failed to create genai client: %w", err)
	}
	modelName := "text-embedding-005"
	// See the documentation: https://pkg.go.dev/google.golang.org/genai#Batches.Create
	job, err := client.Batches.Create(ctx,
		modelName,
		&genai.BatchJobSource{
			Format: "jsonl",
			// Source link: https://storage.cloud.google.com/cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl
			GCSURI: []string{"gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl"},
		},
		&genai.CreateBatchJobConfig{
			Dest: &genai.BatchJobDestination{
				Format: "jsonl",
				GCSURI: outputURI,
			},
		},
	)
	if err != nil {
		return fmt.Errorf("failed to create batch job: %w", err)
	}

	fmt.Fprintf(w, "Job name: %s\n", job.Name)
	fmt.Fprintf(w, "Job state: %s\n", job.State)
	// Example response:
	//  Job name: projects/{PROJECT_ID}/locations/us-central1/batchPredictionJobs/9876453210000000000
	//  Job state: JOB_STATE_PENDING

	// See the documentation: https://pkg.go.dev/google.golang.org/genai#BatchJob
	completedStates := map[genai.JobState]bool{
		genai.JobStateSucceeded: true,
		genai.JobStateFailed:    true,
		genai.JobStateCancelled: true,
		genai.JobStatePaused:    true,
	}

	// Poll until job finishes
	for !completedStates[job.State] {
		time.Sleep(30 * time.Second)
		job, err = client.Batches.Get(ctx, job.Name, nil)
		if err != nil {
			return fmt.Errorf("failed to get batch job: %w", err)
		}
		fmt.Fprintf(w, "Job state: %s\n", job.State)

		if job.State == genai.JobStateFailed {
			fmt.Fprintf(w, "Error: %+v\n", 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

	return nil
}

Node.js

安装

npm install @google/genai

如需了解详情,请参阅 SDK 参考文档

设置环境变量以将 Google 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_ENTERPRISE=True

const {GoogleGenAI} = require('@google/genai');

const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
const GOOGLE_CLOUD_LOCATION =
  process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
const OUTPUT_URI = 'gs://your-bucket/your-prefix';

async function runBatchPredictionJob(
  outputUri = OUTPUT_URI,
  projectId = GOOGLE_CLOUD_PROJECT,
  location = GOOGLE_CLOUD_LOCATION
) {
  const client = new GoogleGenAI({
    vertexai: true,
    project: projectId,
    location: location,
    httpOptions: {
      apiVersion: 'v1',
    },
  });

  // See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
  let job = await client.batches.create({
    model: 'text-embedding-005',
    // Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
    src: 'gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl',
    config: {
      dest: outputUri,
    },
  });

  console.log(`Job name: ${job.name}`);
  console.log(`Job state: ${job.state}`);

  // Example response:
  //  Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
  //  Job state: JOB_STATE_PENDING

  const completedStates = new Set([
    'JOB_STATE_SUCCEEDED',
    'JOB_STATE_FAILED',
    'JOB_STATE_CANCELLED',
    'JOB_STATE_PAUSED',
  ]);

  while (!completedStates.has(job.state)) {
    await new Promise(resolve => setTimeout(resolve, 30000));
    job = await client.batches.get({name: job.name});
    console.log(`Job state: ${job.state}`);
  }

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

  return job.state;
}

Java

了解如何安装或更新 Java

如需了解详情,请参阅 SDK 参考文档

设置环境变量以将 Google 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_ENTERPRISE=True


import static com.google.genai.types.JobState.Known.JOB_STATE_CANCELLED;
import static com.google.genai.types.JobState.Known.JOB_STATE_FAILED;
import static com.google.genai.types.JobState.Known.JOB_STATE_PAUSED;
import static com.google.genai.types.JobState.Known.JOB_STATE_SUCCEEDED;

import com.google.genai.Client;
import com.google.genai.types.BatchJob;
import com.google.genai.types.BatchJobDestination;
import com.google.genai.types.BatchJobSource;
import com.google.genai.types.CreateBatchJobConfig;
import com.google.genai.types.GetBatchJobConfig;
import com.google.genai.types.HttpOptions;
import com.google.genai.types.JobState;
import java.util.EnumSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

public class BatchPredictionEmbeddingsWithGcs {

  public static void main(String[] args) throws InterruptedException {
    // TODO(developer): Replace these variables before running the sample.
    String modelId = "text-embedding-005";
    String outputGcsUri = "gs://your-bucket/your-prefix";
    createBatchJob(modelId, outputGcsUri);
  }

  // Creates a batch prediction job with embedding model and Google Cloud Storage.
  public static JobState createBatchJob(String modelId, String outputGcsUri)
      throws InterruptedException {
    // Client Initialization. Once created, it can be reused for multiple requests.
    try (Client client =
        Client.builder()
            .location("us-central1")
            .vertexAI(true)
            .httpOptions(HttpOptions.builder().apiVersion("v1").build())
            .build()) {

      // See the documentation:
      // https://googleapis.github.io/java-genai/javadoc/com/google/genai/Batches.html
      BatchJobSource batchJobSource =
          BatchJobSource.builder()
              // Source link:
              // https://storage.cloud.google.com/cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl
              .gcsUri("gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl")
              .format("jsonl")
              .build();

      CreateBatchJobConfig batchJobConfig =
          CreateBatchJobConfig.builder()
              .displayName("your-display-name")
              .dest(BatchJobDestination.builder().gcsUri(outputGcsUri).format("jsonl").build())
              .build();

      BatchJob batchJob = client.batches.create(modelId, batchJobSource, batchJobConfig);

      String jobName =
          batchJob.name().orElseThrow(() -> new IllegalStateException("Missing job name"));
      JobState jobState =
          batchJob.state().orElseThrow(() -> new IllegalStateException("Missing job state"));
      System.out.println("Job name: " + jobName);
      System.out.println("Job state: " + jobState);
      // Job name: projects/.../locations/.../batchPredictionJobs/6205497615459549184
      // Job state: JOB_STATE_PENDING

      // See the documentation:
      // https://googleapis.github.io/java-genai/javadoc/com/google/genai/types/BatchJob.html
      Set<JobState.Known> completedStates =
          EnumSet.of(JOB_STATE_SUCCEEDED, JOB_STATE_FAILED, JOB_STATE_CANCELLED, JOB_STATE_PAUSED);

      while (!completedStates.contains(jobState.knownEnum())) {
        TimeUnit.SECONDS.sleep(30);
        batchJob = client.batches.get(jobName, GetBatchJobConfig.builder().build());
        jobState =
            batchJob
                .state()
                .orElseThrow(() -> new IllegalStateException("Missing job state during polling"));
        System.out.println("Job state: " + jobState);
      }
      // Example response:
      // Job state: JOB_STATE_QUEUED
      // Job state: JOB_STATE_RUNNING
      // ...
      // Job state: JOB_STATE_SUCCEEDED
      return jobState;
    }
  }
}

检索批量输出

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

后续步骤