Files
2026-07-13 13:30:30 +08:00

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
}