Receber inferências de embeddings de texto em lote

Receber respostas em lote é uma maneira de enviar com eficiência um grande número de solicitações de embeddings não sensíveis à latência. Diferente da obtenção de respostas on-line, em que há limitação para uma solicitação de entrada por vez, é possível enviar um grande número de solicitações do LLM em uma única solicitação em lote. De maneira semelhante à maneira como a inferência em lote é feita para dados tabulares na Vertex AI, você determina o local de saída, adiciona as entradas e as respostas são preenchidas de maneira assíncrona no local de saída.

Modelos de embeddings de texto compatíveis com inferências em lote

Todas as versões estáveis dos modelos de embedding de texto são compatíveis com inferências em lote, exceto os embeddings do Gemini (gemini-embedding-001). As versões estáveis têm suporte total para ambientes de produção. Para conferir a lista completa de modelos de embedding, consulte Modelo e versões de embedding.

Preparar suas entradas

A entrada para solicitações em lote é uma lista de comandos que podem ser armazenados em uma tabela do BigQuery ou como um arquivo JSON Lines (JSONL) no Cloud Storage. Cada solicitação pode incluir até 30.000 comandos.

Exemplo de JSONL

Esta seção mostra exemplos de como formatar entradas e saídas JSONL.

Exemplo de entrada JSONL

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

Exemplo de saída 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":""}

Exemplo do BigQuery

Esta seção mostra exemplos de como formatar entradas e saídas do BigQuery.

Exemplo de uma entrada do BigQuery

Este exemplo mostra uma tabela do BigQuery de coluna única.

conteúdo
"Faça uma breve descrição de um modelo de machine learning:"
"Melhor receita para pão de banana:"

Exemplo de uma saída do BigQuery

conteúdo	previsões	status
"Faça uma breve descrição de um modelo de machine learning:"	'[{"embeddings": { "statistics":{"token_count":8,"truncated":false}, "Values":[0.1,....] } } ]'
"Melhor receita para pão de banana:"	'[{"embeddings": { "statistics":{"token_count":3,"truncated":false}, "Values":[0.2,....] } } ]'

Solicitar uma resposta em lote

Dependendo do número de itens de entrada enviados, uma tarefa de previsão em lote pode levar algum tempo para ser concluída.

REST

Para testar um prompt de texto usando a API Vertex AI, envie uma solicitação POST para o endpoint do modelo do editor.

Antes de usar os dados da solicitação abaixo, faça as substituições a seguir:

• PROJECT_ID: o ID do seu projeto do Google Cloud .
• BP_JOB_NAME: o nome do job.
• INPUT_URI: o URI da origem da entrada. Pode ser um URI de tabela do BigQuery ou um URI de arquivo JSONL no Cloud Storage.
• OUTPUT_URI: URI de destino de saída.

Método HTTP e URL:

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

Corpo JSON da solicitação:

{
    "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"
    }
  }
}

Para enviar a solicitação, escolha uma destas opções:

curl

Observação: o comando a seguir pressupõe que você fez login na CLI gcloud com sua conta de usuário executando gcloud init ou gcloud auth login, ou usando o Cloud Shell, que faz login automaticamente na CLI gcloud. . É possível verificar a conta ativa atual executando gcloud auth list.

Salve o corpo da solicitação em um arquivo com o nome request.json e execute o comando a seguir:

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

Observação: o comando a seguir pressupõe que você fez login na CLI gcloud com sua conta de usuário executando gcloud init ou gcloud auth login . É possível verificar a conta ativa atual executando gcloud auth list.

Salve o corpo da solicitação em um arquivo com o nome request.json e execute o comando a seguir:

$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

Você receberá uma resposta JSON semelhante a esta:

{
  "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": {}
}

A resposta inclui um identificador exclusivo para a tarefa em lote. É possível pesquisar o status da tarefa em lote usando BATCH_JOB_ID até que o job state seja JOB_STATE_SUCCEEDED. Por exemplo:

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

Observação: só é possível executar um job de resposta em lote por vez. Relatórios de VPC-SC, contas de serviço personalizadas, progresso em tempo real e CMEK e não são compatíveis no momento.

Python

Instalar

pip install --upgrade google-genai

Para saber mais, consulte a documentação de referência do SDK.

Defina variáveis de ambiente para usar o SDK de IA generativa com a 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/.../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

Saiba como instalar ou atualizar o Go.

Para saber mais, consulte a documentação de referência do SDK.

Defina variáveis de ambiente para usar o SDK de IA generativa com a 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 (
	"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

Instalar

npm install @google/genai

Para saber mais, consulte a documentação de referência do SDK.

Defina variáveis de ambiente para usar o SDK de IA generativa com a 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

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

Saiba como instalar ou atualizar o Java.

Para saber mais, consulte a documentação de referência do SDK.

Defina variáveis de ambiente para usar o SDK de IA generativa com a 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 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;
    }
  }
}

Recuperar saída em lote

Quando uma tarefa de inferência em lote é concluída, a saída é armazenada no bucket do Cloud Storage ou na tabela do BigQuery especificada na solicitação.

A seguir

• Saiba como usar embeddings de texto.