バッチ テキスト エンベディング予測を取得する

レスポンスをバッチで取得すると、レイテンシの影響を受けない大量のエンベディング リクエストを効率的に送信できます。レスポンスをオンラインで取得する場合は一度に 1 つの入力リクエストしか行えませんが、バッチで取得する場合は 1 つのバッチ リクエストで多くの LLM リクエストを送信できます。Vertex AI の表形式データに対してバッチ予測を行う場合と同様に、出力先を決めて入力を追加すると、その出力先に非同期でレスポンスが返されます。

バッチ予測をサポートするテキスト エンベディング モデル

テキスト エンベディング モデルのすべての安定版は、textembedding-gecko-multilingual@001 を除き、バッチ予測をサポートしています。安定版はプレビュー版ではなく、本番環境で完全にサポートされているバージョンです。サポートされているエンベディング モデルの一覧については、エンベディング モデルとバージョンをご覧ください。

入力を準備する

バッチ リクエストの入力は、BigQuery テーブルまたは Cloud Storage の JSON Lines（JSONL）ファイルとして保存可能なプロンプトのリストです。1 つのリクエストには最大 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 出力の例

コンテンツ 予測 ステータス
「ML モデルについて簡単に説明してください」 
'[{"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 メソッドと URL:

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

レスポンスには、バッチジョブの固有識別子が含まれます。ジョブ state が JOB_STATE_SUCCEEDED になるまで、BATCH_JOB_ID を使用してバッチジョブのステータスをポーリングできます。例:

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

Vertex AI SDK for Python

Vertex AI SDK for Python のインストールまたは更新の方法については、Vertex AI SDK for Python をインストールするをご覧ください。 詳細については、Vertex AI SDK for Python 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 テーブルに出力が保存されます。

次のステップ