Generative AI on Google Cloud

75 / 80·notebook

Chapter 75

Intro to Batch Inference with the Gemini API

NotebookPython 354 cells

In [ ]python · cell 1

python

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Intro to Batch Inference with the Gemini API


Author(s)	Eric Dong, Holt Skinner

Overview

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 or BigQuery.

Batch inference is generally more efficient and cost-effective than online inference when processing a large number of inputs that are not latency sensitive.

To learn more, see the Batch inference with Gemini page.

Objectives

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.

You will complete the following tasks:

Preparing batch inputs and an output location
Submitting a batch prediction job
Retrieving batch prediction results

Get started

Install libraries

In [ ]python · cell 7

python

%pip install --upgrade --quiet google-genai pandas google-cloud-storage google-cloud-bigquery

⚠️ Note: Ignore pip dependency errors.

Import libraries

In [ ]python · cell 10

python

import os
import sys
import time
from datetime import datetime

import fsspec
import pandas as pd
from google import genai
from google.cloud import bigquery
from google.genai.types import CreateBatchJobConfig

Authenticate your notebook environment (Colab only)

If you're running this notebook on Google Colab, run the cell below to authenticate your environment.

In [ ]python · cell 12

python

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

Autenticate your Google Cloud project

You can use a Google Cloud Project or an API Key for authentication. This tutorial uses a Google Cloud Project.

Enable the Agent Platform API

In [ ]python · cell 14

python

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "global")

In [ ]python · cell 15

python

client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION)

Load model

You can find a list of the Gemini models that support batch inference in the Get batch predictions for Gemini page.

This tutorial uses Gemini 3.5 Flash (gemini-3.5-flash) model.

In [ ]python · cell 17

python

MODEL_ID = "gemini-3.5-flash"  # @param {type:"string", isTemplate: true}

Cloud Storage

Prepare batch inputs

The input for batch requests specifies the items to send to your model for prediction.

This tutorial uses Cloud Storage as an example. The requirements for Cloud Storage input are:

File format: JSON Lines (JSONL)
Multiple files are supported with regex such as gs://bucketname/path/to/*.jsonl
Located in us-central1
Appropriate read permissions for the service account

Each request that you send to a model can include parameters that control how the model generates a response.

This is one of the example requests in the input JSONL file batch_requests_for_multimodal_input_2.jsonl:

json

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

In [ ]python · cell 20

python

# fmt: off
INPUT_DATA = "gs://cloud-samples-data/generative-ai/batch/batch_requests_for_multimodal_input_2.jsonl"  # @param {type:"string"}
# fmt: on

Prepare batch output location

When a batch prediction task completes, the output is stored in the location that you specified in your request.

The location is in the form of a Cloud Storage prefix.
- For example: gs://path/to/output/data.
You can specify the URI of your Cloud Storage bucket in BUCKET_URI, or
If it is not specified, this notebook will create a Cloud Storage bucket in the form of gs://PROJECT_ID-TIMESTAMP.

In [ ]python · cell 22

python

BUCKET_URI = "[your-cloud-storage-bucket]"  # @param {type:"string"}
GCS_LOCATION = "us-central1"  # @param {type:"string"}

if BUCKET_URI == "[your-cloud-storage-bucket]":
    TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")
    BUCKET_URI = f"gs://{PROJECT_ID}-{TIMESTAMP}"

    ! gcloud storage buckets create {BUCKET_URI} --project={PROJECT_ID} --location={GCS_LOCATION}

Send a batch prediction request

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.

To learn more, see the Batch prediction API page.

In [ ]python · cell 24

python

gcs_batch_job = client.batches.create(
    model=MODEL_ID,
    src=INPUT_DATA,
    config=CreateBatchJobConfig(dest=BUCKET_URI),
)
gcs_batch_job.name

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

In [ ]python · cell 26

python

gcs_batch_job = client.batches.get(name=gcs_batch_job.name)
gcs_batch_job

Optionally, you can list all the batch prediction jobs in the project.

In [ ]python · cell 28

python

for job in client.batches.list():
    print(job.name, job.create_time, job.state)

Wait for the batch prediction job to complete

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.

In [ ]python · cell 30

python

# Refresh the job until complete
while gcs_batch_job.state in (
    "JOB_STATE_RUNNING",
    "JOB_STATE_PENDING",
    "JOB_STATE_QUEUED",
):
    time.sleep(5)
    gcs_batch_job = client.batches.get(name=gcs_batch_job.name)

# Check if the job succeeds
if gcs_batch_job.state == "JOB_STATE_SUCCEEDED":
    print("Job succeeded!")
else:
    print(f"Job failed: {gcs_batch_job.error}")

Retrieve batch prediction results

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.

The file name should look like this: {gcs_batch_job.dest.gcs_uri}/prediction-model-TIMESTAMP/predictions.jsonl

Example output:

json

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

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.

You can retrieve the specific responses in the response field.

In [ ]python · cell 33

python

fs = fsspec.filesystem("gcs")

file_paths = fs.glob(f"{gcs_batch_job.dest.gcs_uri}/*/predictions.jsonl")

if gcs_batch_job.state == "JOB_STATE_SUCCEEDED":
    # Load the JSONL file into a DataFrame
    df = pd.read_json(f"gs://{file_paths[0]}", lines=True)

    df = df.join(pd.json_normalize(df["response"], "candidates"))
    display(df)

BigQuery

Batch Input Preparation

To send batch requests for prediction, you need to structure your input properly.

This guide uses BigQuery as an example. To use a BigQuery table as input:

Ensure the dataset is created in a supported region (e.g., us-central1). Multi-region locations (e.g., us) are not allowed.
The input table must include a request column of type JSON or STRING containing valid JSON, structured as a GenerateContentRequest.
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.
Only public YouTube or Cloud Storage URIs are supported in the fileData or file_data field.
Requests can include parameters to customize the model's output.

This is an example BigQuery table with sample requests:

In [ ]python · cell 37

python

# fmt: off
INPUT_DATA = "bq://storage-samples.generative_ai.batch_requests_for_multimodal_input_2"  # @param {type:"string"}
# fmt: on

You can query the BigQuery table to review the input data.

In [ ]python · cell 39

python

bq_client = bigquery.Client(project=PROJECT_ID)

bq_table_id = INPUT_DATA.replace("bq://", "")
sql = f"""
        SELECT *
        FROM {bq_table_id}
        """

query_result = bq_client.query(sql)

df = query_result.result().to_dataframe()
df.head()

Prepare batch output location

When a batch prediction task completes, the output is stored in the location that you specified in your request.

The location is in the form of a BigQuery URI prefix, for example: bq://projectId.bqDatasetId.
If not specified, bq://PROJECT_ID.gen_ai_batch_prediction.predictions_TIMESTAMP will be used.

This tutorial uses a BigQuery table as an example.

You can specify the URI of your BigQuery table in BQ_OUTPUT_URI, or
If it is not specified, this notebook will create a new dataset bq://PROJECT_ID.gen_ai_batch_prediction for you.

In [ ]python · cell 41

python

BQ_OUTPUT_URI = "[your-bigquery-table]"  # @param {type:"string"}

if BQ_OUTPUT_URI == "[your-bigquery-table]":
    bq_dataset_id = "gen_ai_batch_prediction"

    # The output table will be created automatically if it doesn't exist
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    bq_table_id = f"prediction_result_{timestamp}"
    BQ_OUTPUT_URI = f"bq://{PROJECT_ID}.{bq_dataset_id}.{bq_table_id}"

    bq_dataset = bigquery.Dataset(f"{PROJECT_ID}.{bq_dataset_id}")
    bq_dataset.location = "us-central1"

    bq_dataset = bq_client.create_dataset(bq_dataset, exists_ok=True, timeout=30)
    print(
        f"Created BigQuery dataset {bq_client.project}.{bq_dataset.dataset_id} for batch prediction output."
    )

print(f"BigQuery output URI: {BQ_OUTPUT_URI}")

Send a batch prediction request

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.

To learn more, see the Batch prediction API page.

In [ ]python · cell 43

python

bq_batch_job = client.batches.create(
    model=MODEL_ID,
    src=INPUT_DATA,
    config=CreateBatchJobConfig(dest=BQ_OUTPUT_URI),
)
bq_batch_job.name

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

In [ ]python · cell 45

python

bq_batch_job = client.batches.get(name=bq_batch_job.name)
bq_batch_job

Optionally, you can list all the batch prediction jobs in the project.

In [ ]python · cell 47

python

for job in client.batches.list():
    print(job.name, job.create_time, job.state)

Wait for the batch prediction job to complete

In [ ]python · cell 49

python

# Refresh the job until complete
while bq_batch_job.state in (
    "JOB_STATE_RUNNING",
    "JOB_STATE_PENDING",
    "JOB_STATE_QUEUED",
):
    time.sleep(5)
    bq_batch_job = client.batches.get(name=bq_batch_job.name)

# Check if the job succeeds
if bq_batch_job.state == "JOB_STATE_SUCCEEDED":
    print("Job succeeded!")
else:
    print(f"Job failed: {bq_batch_job.error}")

Retrieve batch prediction results

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.

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.
If you did not provide an output dataset, a default dataset bq://PROJECT_ID.gen_ai_batch_prediction will be created for you.
The name of the table is formed by appending predictions_ with the timestamp of when the batch prediction job started.

You can use the example code below to retrieve predictions and store them into a Pandas DataFrame.

In [ ]python · cell 52

python

bq_table_id = bq_batch_job.dest.bigquery_uri.replace("bq://", "")

sql = f"""
        SELECT *
        FROM {bq_table_id}
        """

query_result = bq_client.query(sql)

df = query_result.result().to_dataframe()
df.head()

Cleaning up

Clean up resources created in this notebook.

In [ ]python · cell 54

python

# Delete the batch prediction jobs
if gcs_batch_job:
    client.batches.delete(name=gcs_batch_job.name)
if bq_batch_job:
    client.batches.delete(name=bq_batch_job.name)

← 返回课程

Generative AI on Google Cloud

75 / 80·notebook

Chapter 75

Intro to Batch Inference with the Gemini API

NotebookPython 354 cells

In [ ]python · cell 1

python

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Intro to Batch Inference with the Gemini API


Author(s)	Eric Dong, Holt Skinner

Overview

Batch inference is generally more efficient and cost-effective than online inference when processing a large number of inputs that are not latency sensitive.

To learn more, see the Batch inference with Gemini page.

Objectives

You will complete the following tasks:

Preparing batch inputs and an output location
Submitting a batch prediction job
Retrieving batch prediction results

Get started

Install libraries

In [ ]python · cell 7

python

%pip install --upgrade --quiet google-genai pandas google-cloud-storage google-cloud-bigquery

⚠️ Note: Ignore pip dependency errors.

Import libraries

In [ ]python · cell 10

python

import os
import sys
import time
from datetime import datetime

import fsspec
import pandas as pd
from google import genai
from google.cloud import bigquery
from google.genai.types import CreateBatchJobConfig

Authenticate your notebook environment (Colab only)

If you're running this notebook on Google Colab, run the cell below to authenticate your environment.

In [ ]python · cell 12

python

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

Autenticate your Google Cloud project

You can use a Google Cloud Project or an API Key for authentication. This tutorial uses a Google Cloud Project.

Enable the Agent Platform API

In [ ]python · cell 14

python

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "global")

In [ ]python · cell 15

python

client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION)

Load model

You can find a list of the Gemini models that support batch inference in the Get batch predictions for Gemini page.

This tutorial uses Gemini 3.5 Flash (gemini-3.5-flash) model.

In [ ]python · cell 17

python

MODEL_ID = "gemini-3.5-flash"  # @param {type:"string", isTemplate: true}

Cloud Storage

Prepare batch inputs

The input for batch requests specifies the items to send to your model for prediction.

This tutorial uses Cloud Storage as an example. The requirements for Cloud Storage input are:

File format: JSON Lines (JSONL)
Multiple files are supported with regex such as gs://bucketname/path/to/*.jsonl
Located in us-central1
Appropriate read permissions for the service account

Each request that you send to a model can include parameters that control how the model generates a response.

This is one of the example requests in the input JSONL file batch_requests_for_multimodal_input_2.jsonl:

json

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

In [ ]python · cell 20

python

# fmt: off
INPUT_DATA = "gs://cloud-samples-data/generative-ai/batch/batch_requests_for_multimodal_input_2.jsonl"  # @param {type:"string"}
# fmt: on

Prepare batch output location

When a batch prediction task completes, the output is stored in the location that you specified in your request.

The location is in the form of a Cloud Storage prefix.
- For example: gs://path/to/output/data.
You can specify the URI of your Cloud Storage bucket in BUCKET_URI, or
If it is not specified, this notebook will create a Cloud Storage bucket in the form of gs://PROJECT_ID-TIMESTAMP.

In [ ]python · cell 22

python

BUCKET_URI = "[your-cloud-storage-bucket]"  # @param {type:"string"}
GCS_LOCATION = "us-central1"  # @param {type:"string"}

if BUCKET_URI == "[your-cloud-storage-bucket]":
    TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")
    BUCKET_URI = f"gs://{PROJECT_ID}-{TIMESTAMP}"

    ! gcloud storage buckets create {BUCKET_URI} --project={PROJECT_ID} --location={GCS_LOCATION}

Send a batch prediction request

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.

To learn more, see the Batch prediction API page.

In [ ]python · cell 24

python

gcs_batch_job = client.batches.create(
    model=MODEL_ID,
    src=INPUT_DATA,
    config=CreateBatchJobConfig(dest=BUCKET_URI),
)
gcs_batch_job.name

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

In [ ]python · cell 26

python

gcs_batch_job = client.batches.get(name=gcs_batch_job.name)
gcs_batch_job

Optionally, you can list all the batch prediction jobs in the project.

In [ ]python · cell 28

python

for job in client.batches.list():
    print(job.name, job.create_time, job.state)

Wait for the batch prediction job to complete

In [ ]python · cell 30

python

# Refresh the job until complete
while gcs_batch_job.state in (
    "JOB_STATE_RUNNING",
    "JOB_STATE_PENDING",
    "JOB_STATE_QUEUED",
):
    time.sleep(5)
    gcs_batch_job = client.batches.get(name=gcs_batch_job.name)

# Check if the job succeeds
if gcs_batch_job.state == "JOB_STATE_SUCCEEDED":
    print("Job succeeded!")
else:
    print(f"Job failed: {gcs_batch_job.error}")

Retrieve batch prediction results

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.

The file name should look like this: {gcs_batch_job.dest.gcs_uri}/prediction-model-TIMESTAMP/predictions.jsonl

Example output:

json

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

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.

You can retrieve the specific responses in the response field.

In [ ]python · cell 33

python

fs = fsspec.filesystem("gcs")

file_paths = fs.glob(f"{gcs_batch_job.dest.gcs_uri}/*/predictions.jsonl")

if gcs_batch_job.state == "JOB_STATE_SUCCEEDED":
    # Load the JSONL file into a DataFrame
    df = pd.read_json(f"gs://{file_paths[0]}", lines=True)

    df = df.join(pd.json_normalize(df["response"], "candidates"))
    display(df)

BigQuery

Batch Input Preparation

To send batch requests for prediction, you need to structure your input properly.

This guide uses BigQuery as an example. To use a BigQuery table as input:

Ensure the dataset is created in a supported region (e.g., us-central1). Multi-region locations (e.g., us) are not allowed.
The input table must include a request column of type JSON or STRING containing valid JSON, structured as a GenerateContentRequest.
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.
Only public YouTube or Cloud Storage URIs are supported in the fileData or file_data field.
Requests can include parameters to customize the model's output.

This is an example BigQuery table with sample requests:

In [ ]python · cell 37

python

# fmt: off
INPUT_DATA = "bq://storage-samples.generative_ai.batch_requests_for_multimodal_input_2"  # @param {type:"string"}
# fmt: on

You can query the BigQuery table to review the input data.

In [ ]python · cell 39

python

bq_client = bigquery.Client(project=PROJECT_ID)

bq_table_id = INPUT_DATA.replace("bq://", "")
sql = f"""
        SELECT *
        FROM {bq_table_id}
        """

query_result = bq_client.query(sql)

df = query_result.result().to_dataframe()
df.head()

Prepare batch output location

When a batch prediction task completes, the output is stored in the location that you specified in your request.

The location is in the form of a BigQuery URI prefix, for example: bq://projectId.bqDatasetId.
If not specified, bq://PROJECT_ID.gen_ai_batch_prediction.predictions_TIMESTAMP will be used.

This tutorial uses a BigQuery table as an example.

You can specify the URI of your BigQuery table in BQ_OUTPUT_URI, or
If it is not specified, this notebook will create a new dataset bq://PROJECT_ID.gen_ai_batch_prediction for you.

In [ ]python · cell 41

python

BQ_OUTPUT_URI = "[your-bigquery-table]"  # @param {type:"string"}

if BQ_OUTPUT_URI == "[your-bigquery-table]":
    bq_dataset_id = "gen_ai_batch_prediction"

    # The output table will be created automatically if it doesn't exist
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    bq_table_id = f"prediction_result_{timestamp}"
    BQ_OUTPUT_URI = f"bq://{PROJECT_ID}.{bq_dataset_id}.{bq_table_id}"

    bq_dataset = bigquery.Dataset(f"{PROJECT_ID}.{bq_dataset_id}")
    bq_dataset.location = "us-central1"

    bq_dataset = bq_client.create_dataset(bq_dataset, exists_ok=True, timeout=30)
    print(
        f"Created BigQuery dataset {bq_client.project}.{bq_dataset.dataset_id} for batch prediction output."
    )

print(f"BigQuery output URI: {BQ_OUTPUT_URI}")

Send a batch prediction request

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.

To learn more, see the Batch prediction API page.

In [ ]python · cell 43

python

bq_batch_job = client.batches.create(
    model=MODEL_ID,
    src=INPUT_DATA,
    config=CreateBatchJobConfig(dest=BQ_OUTPUT_URI),
)
bq_batch_job.name

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

In [ ]python · cell 45

python

bq_batch_job = client.batches.get(name=bq_batch_job.name)
bq_batch_job

Optionally, you can list all the batch prediction jobs in the project.

In [ ]python · cell 47

python

for job in client.batches.list():
    print(job.name, job.create_time, job.state)

Wait for the batch prediction job to complete

In [ ]python · cell 49

python

# Refresh the job until complete
while bq_batch_job.state in (
    "JOB_STATE_RUNNING",
    "JOB_STATE_PENDING",
    "JOB_STATE_QUEUED",
):
    time.sleep(5)
    bq_batch_job = client.batches.get(name=bq_batch_job.name)

# Check if the job succeeds
if bq_batch_job.state == "JOB_STATE_SUCCEEDED":
    print("Job succeeded!")
else:
    print(f"Job failed: {bq_batch_job.error}")

Retrieve batch prediction results

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.
If you did not provide an output dataset, a default dataset bq://PROJECT_ID.gen_ai_batch_prediction will be created for you.
The name of the table is formed by appending predictions_ with the timestamp of when the batch prediction job started.

You can use the example code below to retrieve predictions and store them into a Pandas DataFrame.

In [ ]python · cell 52

python

bq_table_id = bq_batch_job.dest.bigquery_uri.replace("bq://", "")

sql = f"""
        SELECT *
        FROM {bq_table_id}
        """

query_result = bq_client.query(sql)

df = query_result.result().to_dataframe()
df.head()

Cleaning up

Clean up resources created in this notebook.

In [ ]python · cell 54

python

# Delete the batch prediction jobs
if gcs_batch_job:
    client.batches.delete(name=gcs_batch_job.name)
if bq_batch_job:
    client.batches.delete(name=bq_batch_job.name)