Obtén predicciones de incorporaciones de texto por lotes

Obtener respuestas en un lote es una manera de enviar de manera eficiente grandes cantidades de solicitudes de incorporaciones no sensibles a la latencia. A diferencia de obtener respuestas en línea, en la que estás limitado a una solicitud de entrada a la vez, puedes enviar una gran cantidad de solicitudes de LLM en una sola solicitud por lotes. De manera similar a como se realiza la predicción por lotes para los datos tabulares en Vertex AI, debes determinar la ubicación de la salida, agregar la entrada y las respuestas se propagan de forma asíncrona en la ubicación de salida.

Modelos de incorporación de texto que admiten predicciones por lotes

Todas las versiones estables del modelo textembedding-gecko admiten predicciones por lotes, excepto para textembedding-gecko-multilingual@001. Las versiones estables son versiones que ya no se encuentran en la vista previa y que son totalmente compatibles con la entornos de producción. Para ver la lista completa de modelos de incorporación compatibles, consulta Modelo y versiones de incorporación.

Prepara tus entradas

La entrada para las solicitudes por lotes es una lista de mensajes que se pueden almacenar en una tabla de BigQuery o como un archivo de líneas JSON (JSONL) en Cloud Storage. Cada solicitud puede incluir hasta 30,000 mensajes.

Ejemplo de JSONL

En esta sección, se muestran ejemplos de cómo dar formato a la entrada y la salida de JSONL.

Ejemplo de entrada JSONL

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

Ejemplo de salida 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":""}

Ejemplo de BigQuery

En esta sección, se muestran ejemplos de cómo dar formato a la entrada y la salida de BigQuery.

Ejemplo de entrada de BigQuery

En este ejemplo, se muestra una tabla de BigQuery de una sola columna.

contenido
“Proporciona una breve descripción de un modelo de aprendizaje automático:”
“Mejor receta de pan de banana”.

Ejemplo de salida de BigQuery

contenido predicciones estado
“Proporciona una breve descripción de un modelo de aprendizaje automático:” 
'[{"embeddings":
    { "statistics":{"token_count":8,"truncated":false},
      "Values":[0.1,....]
    }
  }
]'
“Mejor receta de pan de banana”. 
'[{"embeddings":
    { "statistics":{"token_count":3,"truncated":false},
      "Values":[0.2,....]
    }
  }
]'

Solicita una respuesta por lotes

Según la cantidad de elementos de entrada que hayas enviado, una tarea de predicción por lotes puede tardar un tiempo en completarse.

REST

Para probar un mensaje de texto con la API de Vertex AI, envía una solicitud POST al extremo del modelo de publicador.

Antes de usar cualquiera de los datos de solicitud a continuación, realiza los siguientes reemplazos:

  • PROJECT_ID: El ID del proyecto de Google Cloud.
  • BP_JOB_NAME: El nombre del trabajo.
  • INPUT_URI: Es el URI de la fuente de entrada. Este es un URI de tabla de BigQuery o un URI de archivo JSONL en Cloud Storage.
  • OUTPUT_URI: URI de destino de salida.

Método HTTP y URL: 

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

Cuerpo JSON de la solicitud:

{
    "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 tu solicitud, elige una de estas opciones:

curl

Guarda el cuerpo de la solicitud en un archivo llamado request.json y ejecuta el siguiente comando:

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

Guarda el cuerpo de la solicitud en un archivo llamado request.json y ejecuta el siguiente comando:

$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

Deberías recibir una respuesta JSON similar a la que se muestra a continuación:

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

En la respuesta, se incluye un identificador único para el trabajo por lotes. Puedes consultar el estado del trabajo por lotes mediante BATCH_JOB_ID hasta que el state sea JOB_STATE_SUCCEEDED. Por ejemplo:

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

Python

Si deseas obtener información para instalar o actualizar el SDK de Vertex AI para Python, consulta Instala el SDK de Vertex AI para Python. Si deseas obtener más información, consulta la documentación de referencia de la API de Python. 

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

Recuperar los resultados por lotes

Cuando se completa una tarea de predicción por lotes, el resultado se almacena en el bucket de Cloud Storage o en la tabla de BigQuery que especificaste en la solicitud.

¿Qué sigue?