836 lines
30 KiB
Plaintext
836 lines
30 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ur8xi4C7S06n"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright 2024 Google LLC\n",
|
|
"#\n",
|
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
|
"# you may not use this file except in compliance with the License.\n",
|
|
"# You may obtain a copy of the License at\n",
|
|
"#\n",
|
|
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
|
"#\n",
|
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
|
"# See the License for the specific language governing permissions and\n",
|
|
"# limitations under the License."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "JAPoU8Sm5E6e"
|
|
},
|
|
"source": [
|
|
"# Intro to Batch Inference with the Gemini API\n",
|
|
"\n",
|
|
"\n",
|
|
"<table align=\"left\">\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/agent-platform/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fbatch-prediction%2Fintro_batch_prediction.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/agent-platform/workbench/instances?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/workbench-icon.svg\" alt=\"Workbench logo\"><br> Open in Workbench\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
"</table>\n",
|
|
"\n",
|
|
"<div style=\"clear: both;\"></div>\n",
|
|
"\n",
|
|
"<p>\n",
|
|
"<b>Share to:</b>\n",
|
|
"\n",
|
|
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
|
|
"</a>\n",
|
|
"</p>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "84f0f73a0f76"
|
|
},
|
|
"source": [
|
|
"| | |\n",
|
|
"|-|-|\n",
|
|
"|Author(s) | [Eric Dong](https://github.com/gericdong), [Holt Skinner](https://github.com/holtskinner) |"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "tvgnzT1CKxrO"
|
|
},
|
|
"source": [
|
|
"## Overview\n",
|
|
"\n",
|
|
"Different from getting online (synchronous) responses, where you are limited to one input request at a time, the batch inference with the Gemini API in Agent Platform allow you to send a large number of multimodal requests to a Gemini model in a single batch request. Then, the model responses asynchronously populate to your storage output location in [Cloud Storage](https://cloud.google.com/storage/docs/introduction) or [BigQuery](https://cloud.google.com/bigquery/docs/storage_overview).\n",
|
|
"\n",
|
|
"Batch inference is generally more efficient and cost-effective than online inference when processing a large number of inputs that are not latency sensitive.\n",
|
|
"\n",
|
|
"To learn more, see the [Batch inference with Gemini](https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/capabilities/batch-prediction-gemini) page.\n",
|
|
"\n",
|
|
"### Objectives\n",
|
|
"\n",
|
|
"In this tutorial, you learn how to make batch inference with the Gemini API in Gemini Enterprise Agent Platform. This tutorial shows how to use **Cloud Storage** and **BigQuery** as input sources and output locations.\n",
|
|
"\n",
|
|
"You will complete the following tasks:\n",
|
|
"\n",
|
|
"- Preparing batch inputs and an output location\n",
|
|
"- Submitting a batch prediction job\n",
|
|
"- Retrieving batch prediction results\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "61RBz8LLbxCR"
|
|
},
|
|
"source": [
|
|
"## Get started"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "No17Cw5hgx12"
|
|
},
|
|
"source": [
|
|
"### Install libraries\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "tFy3H3aPgx12"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install --upgrade --quiet google-genai pandas google-cloud-storage google-cloud-bigquery"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "3T65_kungUTt"
|
|
},
|
|
"source": [
|
|
"⚠️ **Note**: Ignore pip dependency errors."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "06489bd14f16"
|
|
},
|
|
"source": [
|
|
"### Import libraries\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "Uyf7QLsvf_Hp"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"import sys\n",
|
|
"import time\n",
|
|
"from datetime import datetime\n",
|
|
"\n",
|
|
"import fsspec\n",
|
|
"import pandas as pd\n",
|
|
"from google import genai\n",
|
|
"from google.cloud import bigquery\n",
|
|
"from google.genai.types import CreateBatchJobConfig"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dmWOrTJ3gx13"
|
|
},
|
|
"source": [
|
|
"### Authenticate your notebook environment (Colab only)\n",
|
|
"\n",
|
|
"If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "NyKGtVQjgx13"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"if \"google.colab\" in sys.modules:\n",
|
|
" from google.colab import auth\n",
|
|
"\n",
|
|
" auth.authenticate_user()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "DF4l8DTdWgPY"
|
|
},
|
|
"source": [
|
|
"### Autenticate your Google Cloud project\n",
|
|
"\n",
|
|
"You can use a Google Cloud Project or an API Key for authentication. This tutorial uses a Google Cloud Project.\n",
|
|
"\n",
|
|
"- [Enable the Agent Platform API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "Nqwi-5ufWp_B"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# fmt: off\n",
|
|
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
|
|
"# fmt: on\n",
|
|
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
|
|
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
|
|
"\n",
|
|
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"global\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "cfca4d7bd6db"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "e43229f3ad4f"
|
|
},
|
|
"source": [
|
|
"### Load model\n",
|
|
"\n",
|
|
"You can find a list of the Gemini models that support batch inference in the [Get batch predictions for Gemini](https://docs.cloud.google.com/gemini-enterprise-agent-platform/reference/models/batch-prediction-api#supported-models) page.\n",
|
|
"\n",
|
|
"This tutorial uses Gemini 3.5 Flash (`gemini-3.5-flash`) model."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "cf93d5f0ce00"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"MODEL_ID = \"gemini-3.5-flash\" # @param {type:\"string\", isTemplate: true}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "265f180b58e0"
|
|
},
|
|
"source": [
|
|
"## Cloud Storage"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "1_xZADsak23H"
|
|
},
|
|
"source": [
|
|
"### Prepare batch inputs\n",
|
|
"\n",
|
|
"The input for batch requests specifies the items to send to your model for prediction.\n",
|
|
"\n",
|
|
"This tutorial uses Cloud Storage as an example. The requirements for Cloud Storage input are:\n",
|
|
"\n",
|
|
"- File format: [JSON Lines (JSONL)](https://jsonlines.org/)\n",
|
|
"- Multiple files are supported with regex such as gs://bucketname/path/to/*.jsonl\n",
|
|
"- Located in `us-central1`\n",
|
|
"- Appropriate read permissions for the service account\n",
|
|
"\n",
|
|
"Each request that you send to a model can include parameters that control how the model generates a response.\n",
|
|
"\n",
|
|
"This is one of the example requests in the input JSONL file `batch_requests_for_multimodal_input_2.jsonl`:\n",
|
|
"\n",
|
|
"```json\n",
|
|
"{\"request\":{\"contents\": [{\"role\": \"user\", \"parts\": [{\"text\": \"List objects in this image.\"}, {\"file_data\": {\"file_uri\": \"gs://cloud-samples-data/generative-ai/image/office-desk.jpeg\", \"mime_type\": \"image/jpeg\"}}]}],\"generationConfig\":{\"temperature\": 0.4}}}\n",
|
|
"```"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "uWb8QzxwbH6W"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# fmt: off\n",
|
|
"INPUT_DATA = \"gs://cloud-samples-data/generative-ai/batch/batch_requests_for_multimodal_input_2.jsonl\" # @param {type:\"string\"}\n",
|
|
"# fmt: on"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "T3jQ59mCsXLc"
|
|
},
|
|
"source": [
|
|
"### Prepare batch output location\n",
|
|
"\n",
|
|
"When a batch prediction task completes, the output is stored in the location that you specified in your request.\n",
|
|
"\n",
|
|
"- The location is in the form of a Cloud Storage prefix.\n",
|
|
" - For example: `gs://path/to/output/data`.\n",
|
|
"\n",
|
|
"- You can specify the URI of your Cloud Storage bucket in `BUCKET_URI`, or\n",
|
|
"- If it is not specified, this notebook will create a Cloud Storage bucket in the form of `gs://PROJECT_ID-TIMESTAMP`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "OtUodwGXZ7US"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"BUCKET_URI = \"[your-cloud-storage-bucket]\" # @param {type:\"string\"}\n",
|
|
"GCS_LOCATION = \"us-central1\" # @param {type:\"string\"}\n",
|
|
"\n",
|
|
"if BUCKET_URI == \"[your-cloud-storage-bucket]\":\n",
|
|
" TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
|
|
" BUCKET_URI = f\"gs://{PROJECT_ID}-{TIMESTAMP}\"\n",
|
|
"\n",
|
|
" ! gcloud storage buckets create {BUCKET_URI} --project={PROJECT_ID} --location={GCS_LOCATION}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "T90CwWDHvonn"
|
|
},
|
|
"source": [
|
|
"### Send a batch prediction request\n",
|
|
"\n",
|
|
"To make a batch prediction request, you specify a source model ID, an input source and an output location where Agent Platform stores the batch prediction results.\n",
|
|
"\n",
|
|
"To learn more, see the [Batch prediction API](https://docs.cloud.google.com/gemini-enterprise-agent-platform/reference/models/batch-prediction-api) page.\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "d8e54c57072e"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"gcs_batch_job = client.batches.create(\n",
|
|
" model=MODEL_ID,\n",
|
|
" src=INPUT_DATA,\n",
|
|
" config=CreateBatchJobConfig(dest=BUCKET_URI),\n",
|
|
")\n",
|
|
"gcs_batch_job.name"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "A-Fo_Kd9FYRj"
|
|
},
|
|
"source": [
|
|
"Print out the job status and other properties. You can also check the status in the Cloud Console at https://console.cloud.google.com/agent-platform/batch-predictions"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "DWq7m79PbjG8"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"gcs_batch_job = client.batches.get(name=gcs_batch_job.name)\n",
|
|
"gcs_batch_job"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "WHUUEREoiewD"
|
|
},
|
|
"source": [
|
|
"Optionally, you can list all the batch prediction jobs in the project."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "QVgOnasfigx1"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"for job in client.batches.list():\n",
|
|
" print(job.name, job.create_time, job.state)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "7aJaPNBrGPqK"
|
|
},
|
|
"source": [
|
|
"### Wait for the batch prediction job to complete\n",
|
|
"\n",
|
|
"Depending on the number of input items that you submitted, a batch generation task can take some time to complete. You can use the following code to check the job status and wait for the job to complete."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "dtJDIXdHc0W-"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Refresh the job until complete\n",
|
|
"while gcs_batch_job.state in (\n",
|
|
" \"JOB_STATE_RUNNING\",\n",
|
|
" \"JOB_STATE_PENDING\",\n",
|
|
" \"JOB_STATE_QUEUED\",\n",
|
|
"):\n",
|
|
" time.sleep(5)\n",
|
|
" gcs_batch_job = client.batches.get(name=gcs_batch_job.name)\n",
|
|
"\n",
|
|
"# Check if the job succeeds\n",
|
|
"if gcs_batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
|
|
" print(\"Job succeeded!\")\n",
|
|
"else:\n",
|
|
" print(f\"Job failed: {gcs_batch_job.error}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "XWUgAxL-HjN9"
|
|
},
|
|
"source": [
|
|
"### Retrieve batch prediction results\n",
|
|
"\n",
|
|
"When a batch prediction task is complete, the output of the prediction is stored in the bucket in JSONL that you specified in your request.\n",
|
|
"\n",
|
|
"The file name should look like this: `{gcs_batch_job.dest.gcs_uri}/prediction-model-TIMESTAMP/predictions.jsonl`\n",
|
|
"\n",
|
|
"Example output:\n",
|
|
"\n",
|
|
"```json\n",
|
|
"{\"status\": \"\", \"processed_time\": \"2024-11-13T14:04:28.376+00:00\", \"request\": {\"contents\": [{\"parts\": [{\"file_data\": null, \"text\": \"List objects in this image.\"}, {\"file_data\": {\"file_uri\": \"gs://cloud-samples-data/generative-ai/image/gardening-tools.jpeg\", \"mime_type\": \"image/jpeg\"}, \"text\": null}], \"role\": \"user\"}], \"generationConfig\": {\"temperature\": 0.4}}, \"response\": {\"candidates\": [{\"avgLogprobs\": -0.10394711927934126, \"content\": {\"parts\": [{\"text\": \"Here's a list of the objects in the image:\\n\\n* **Watering can:** A green plastic watering can with a white rose head.\\n* **Plant:** A small plant (possibly oregano) in a terracotta pot.\\n* **Terracotta pots:** Two terracotta pots, one containing the plant and another empty, stacked on top of each other.\\n* **Gardening gloves:** A pair of striped gardening gloves.\\n* **Gardening tools:** A small trowel and a hand cultivator (hoe). Both are green with black handles.\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\"}], \"modelVersion\": \"gemini-3.5-flash@default\", \"usageMetadata\": {\"candidatesTokenCount\": 110, \"promptTokenCount\": 264, \"totalTokenCount\": 374}}}\n",
|
|
"```\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "qZlRIYsC01F1"
|
|
},
|
|
"source": [
|
|
"The example code below shows how to load the `.jsonl` file in the Cloud Storage output location into a Pandas DataFrame and print out the object.\n",
|
|
"\n",
|
|
"You can retrieve the specific responses in the `response` field."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "-jLl3es3dTqB"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"fs = fsspec.filesystem(\"gcs\")\n",
|
|
"\n",
|
|
"file_paths = fs.glob(f\"{gcs_batch_job.dest.gcs_uri}/*/predictions.jsonl\")\n",
|
|
"\n",
|
|
"if gcs_batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
|
|
" # Load the JSONL file into a DataFrame\n",
|
|
" df = pd.read_json(f\"gs://{file_paths[0]}\", lines=True)\n",
|
|
"\n",
|
|
" df = df.join(pd.json_normalize(df[\"response\"], \"candidates\"))\n",
|
|
" display(df)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "bfb2a462a7c6"
|
|
},
|
|
"source": [
|
|
"## BigQuery"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "5ea69e283023"
|
|
},
|
|
"source": [
|
|
"### Batch Input Preparation \n",
|
|
"\n",
|
|
"To send batch requests for prediction, you need to structure your input properly. \n",
|
|
"\n",
|
|
"This guide uses **BigQuery** as an example. To use a BigQuery table as input: \n",
|
|
"- Ensure the dataset is created in a supported region (e.g., `us-central1`). Multi-region locations (e.g., `us`) are not allowed. \n",
|
|
"- The input table must include a `request` column of type `JSON` or `STRING` containing valid JSON, structured as a `GenerateContentRequest`. \n",
|
|
"- Additional columns can use any BigQuery data types except `array`, `struct`, `range`, `datetime`, and `geography`. These are ignored for generation but appear in the output table. The system reserves `response` and `status` for output. \n",
|
|
"- Only public YouTube or Cloud Storage URIs are supported in the `fileData` or `file_data` field. \n",
|
|
"- Requests can include parameters to customize the model's output."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "835bfd485d49"
|
|
},
|
|
"source": [
|
|
"This is an example BigQuery table with sample requests:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "b932f02d1f9c"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# fmt: off\n",
|
|
"INPUT_DATA = \"bq://storage-samples.generative_ai.batch_requests_for_multimodal_input_2\" # @param {type:\"string\"}\n",
|
|
"# fmt: on"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "7d62eb771ad6"
|
|
},
|
|
"source": [
|
|
"You can query the BigQuery table to review the input data."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "fbfcefb7d295"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"bq_client = bigquery.Client(project=PROJECT_ID)\n",
|
|
"\n",
|
|
"bq_table_id = INPUT_DATA.replace(\"bq://\", \"\")\n",
|
|
"sql = f\"\"\"\n",
|
|
" SELECT *\n",
|
|
" FROM {bq_table_id}\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
"query_result = bq_client.query(sql)\n",
|
|
"\n",
|
|
"df = query_result.result().to_dataframe()\n",
|
|
"df.head()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "529dae543bfc"
|
|
},
|
|
"source": [
|
|
"### Prepare batch output location\n",
|
|
"\n",
|
|
"When a batch prediction task completes, the output is stored in the location that you specified in your request.\n",
|
|
"\n",
|
|
"- The location is in the form of a BigQuery URI prefix, for example: `bq://projectId.bqDatasetId`.\n",
|
|
"- If not specified, `bq://PROJECT_ID.gen_ai_batch_prediction.predictions_TIMESTAMP` will be used.\n",
|
|
"\n",
|
|
"This tutorial uses a **BigQuery** table as an example.\n",
|
|
"\n",
|
|
"- You can specify the URI of your BigQuery table in `BQ_OUTPUT_URI`, or\n",
|
|
"- If it is not specified, this notebook will create a new dataset `bq://PROJECT_ID.gen_ai_batch_prediction` for you."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "75914840e555"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"BQ_OUTPUT_URI = \"[your-bigquery-table]\" # @param {type:\"string\"}\n",
|
|
"\n",
|
|
"if BQ_OUTPUT_URI == \"[your-bigquery-table]\":\n",
|
|
" bq_dataset_id = \"gen_ai_batch_prediction\"\n",
|
|
"\n",
|
|
" # The output table will be created automatically if it doesn't exist\n",
|
|
" timestamp = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
|
|
" bq_table_id = f\"prediction_result_{timestamp}\"\n",
|
|
" BQ_OUTPUT_URI = f\"bq://{PROJECT_ID}.{bq_dataset_id}.{bq_table_id}\"\n",
|
|
"\n",
|
|
" bq_dataset = bigquery.Dataset(f\"{PROJECT_ID}.{bq_dataset_id}\")\n",
|
|
" bq_dataset.location = \"us-central1\"\n",
|
|
"\n",
|
|
" bq_dataset = bq_client.create_dataset(bq_dataset, exists_ok=True, timeout=30)\n",
|
|
" print(\n",
|
|
" f\"Created BigQuery dataset {bq_client.project}.{bq_dataset.dataset_id} for batch prediction output.\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"print(f\"BigQuery output URI: {BQ_OUTPUT_URI}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "14f84475dc26"
|
|
},
|
|
"source": [
|
|
"### Send a batch prediction request\n",
|
|
"\n",
|
|
"To make a batch prediction request, you specify a source model ID, an input source and an output location where Agent Platform stores the batch prediction results.\n",
|
|
"\n",
|
|
"To learn more, see the [Batch prediction API](https://docs.cloud.google.com/gemini-enterprise-agent-platform/reference/models/batch-prediction-api) page.\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "e809955e753b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"bq_batch_job = client.batches.create(\n",
|
|
" model=MODEL_ID,\n",
|
|
" src=INPUT_DATA,\n",
|
|
" config=CreateBatchJobConfig(dest=BQ_OUTPUT_URI),\n",
|
|
")\n",
|
|
"bq_batch_job.name"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "5fcb3dc2a5fc"
|
|
},
|
|
"source": [
|
|
"Print out the job status and other properties. You can also check the status in the Cloud Console at https://console.cloud.google.com/agent-platform/batch-predictions"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "b19319d92bf0"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"bq_batch_job = client.batches.get(name=bq_batch_job.name)\n",
|
|
"bq_batch_job"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "e72800144910"
|
|
},
|
|
"source": [
|
|
"Optionally, you can list all the batch prediction jobs in the project."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "6a9fb087f9ba"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"for job in client.batches.list():\n",
|
|
" print(job.name, job.create_time, job.state)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "ebcb00b3add9"
|
|
},
|
|
"source": [
|
|
"### Wait for the batch prediction job to complete\n",
|
|
"\n",
|
|
"Depending on the number of input items that you submitted, a batch generation task can take some time to complete. You can use the following code to check the job status and wait for the job to complete."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "189945468c6b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Refresh the job until complete\n",
|
|
"while bq_batch_job.state in (\n",
|
|
" \"JOB_STATE_RUNNING\",\n",
|
|
" \"JOB_STATE_PENDING\",\n",
|
|
" \"JOB_STATE_QUEUED\",\n",
|
|
"):\n",
|
|
" time.sleep(5)\n",
|
|
" bq_batch_job = client.batches.get(name=bq_batch_job.name)\n",
|
|
"\n",
|
|
"# Check if the job succeeds\n",
|
|
"if bq_batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
|
|
" print(\"Job succeeded!\")\n",
|
|
"else:\n",
|
|
" print(f\"Job failed: {bq_batch_job.error}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "ec25f40f0dd9"
|
|
},
|
|
"source": [
|
|
"### Retrieve batch prediction results\n",
|
|
"\n",
|
|
"When a batch prediction task is complete, the output of the prediction is stored in the location that you specified in your request. It is also available in `batch_job.dest.bigquery_uri` or `batch_job.dest.gcs_uri`.\n",
|
|
"\n",
|
|
"- When you are using BigQuery, the output of batch prediction is stored in an output dataset. If you had provided a dataset, the name of the dataset (`BQ_OUTPUT_URI`) is the name you had provided earlier.\n",
|
|
"- If you did not provide an output dataset, a default dataset `bq://PROJECT_ID.gen_ai_batch_prediction` will be created for you.\n",
|
|
"- The name of the table is formed by appending `predictions_` with the timestamp of when the batch prediction job started."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "c459121f0169"
|
|
},
|
|
"source": [
|
|
"You can use the example code below to retrieve predictions and store them into a Pandas DataFrame.\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ca8b73b0ad1b"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"bq_table_id = bq_batch_job.dest.bigquery_uri.replace(\"bq://\", \"\")\n",
|
|
"\n",
|
|
"sql = f\"\"\"\n",
|
|
" SELECT *\n",
|
|
" FROM {bq_table_id}\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
"query_result = bq_client.query(sql)\n",
|
|
"\n",
|
|
"df = query_result.result().to_dataframe()\n",
|
|
"df.head()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2a4e033321ad"
|
|
},
|
|
"source": [
|
|
"## Cleaning up\n",
|
|
"\n",
|
|
"Clean up resources created in this notebook."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ZNCyIKIrdPJY"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Delete the batch prediction jobs\n",
|
|
"if gcs_batch_job:\n",
|
|
" client.batches.delete(name=gcs_batch_job.name)\n",
|
|
"if bq_batch_job:\n",
|
|
" client.batches.delete(name=bq_batch_job.name)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "intro_batch_prediction.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|