{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_TI7YxcMPkTi"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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": "ysLSgHUz9JAN"
      },
      "source": [
        "# Intro to Vertex AI Multimodal Datasets\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/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fmultimodal-dataset%2Fintro_vertex_ai_multimodal_dataset.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/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/multimodal-dataset/intro_vertex_ai_multimodal_dataset.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/multimodal-dataset/intro_vertex_ai_multimodal_dataset.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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",
        "<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/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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/multimodal-dataset/intro_vertex_ai_multimodal_dataset.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AIGJwpvY3ELC"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Frances Thoma](https://github.com/diskontinuum) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7veA0yodoHzx"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use Vertex AI Multimodal Datasets to assemble Gemini requests, to run a validation and resource estimation for supervised fine-tuning, and to create tuning and batch prediction jobs.\n",
        "\n",
        "### Objectives\n",
        "\n",
        "- Preview the new Vertex AI Multimodal Datasets SDK\n",
        "- Demo upcoming integrations\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "* BigQuery\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and [BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n",
        "\n",
        "### Prerequisites\n",
        "1. Make sure that [billing is enabled](https://cloud.google.com/billing/docs/how-to/modify-project) for your project.\n",
        "\n",
        "2. You must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n",
        "\n",
        "### Questions or Feedback\n",
        "\n",
        "You can reach out directly to the authors via `vertex-multimodal-dataset-external-feedback@google.com` for feedback or questions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t-49KwqiQiaw"
      },
      "source": [
        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U3fA_vcl3Y8s"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Md24tB4pYTSY"
      },
      "outputs": [],
      "source": [
        "%pip install --quiet --upgrade google-cloud-aiplatform bigframes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DbqCDqs5QlGh"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gx-AV_VkafnO"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dllbnne2al7X"
      },
      "source": [
        "- If you are running this notebook in a local development environment:\n",
        "  - Install the [Google Cloud SDK](https://cloud.google.com/sdk).\n",
        "  - Obtain authentication credentials. Create local credentials by running the following command and following the oauth2 flow (read more about the command [here](https://cloud.google.com/sdk/gcloud/reference/beta/auth/application-default/login)):\n",
        "\n",
        "    ```bash\n",
        "    gcloud auth application-default login\n",
        "    ```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-xkKPNED9yPa"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mp_q41d8RgAr"
      },
      "outputs": [],
      "source": [
        "import io\n",
        "import os\n",
        "\n",
        "import bigframes.pandas as bpd\n",
        "import pandas\n",
        "import vertexai\n",
        "from PIL import Image\n",
        "from google.cloud import storage\n",
        "from google.cloud.aiplatform.preview import datasets\n",
        "from vertexai.generative_models import Content, Part"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9OCHaQsliNJ"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HptYirwARjXW"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "# 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\", \"us-central1\")\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "# BigFrames settings\n",
        "bpd.close_session()\n",
        "bpd.options.bigquery.project = PROJECT_ID\n",
        "bpd.options.bigquery.location = LOCATION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72tXYePAWDGJ"
      },
      "source": [
        "### Data preparation\n",
        "\n",
        "The image files and labels used in this tutorial are from the flower dataset used in this [TensorFlow blog post](https://cloud.google.com/blog/products/gcp/how-to-classify-images-with-tensorflow-using-google-cloud-machine-learning-and-cloud-dataflow).\n",
        "\n",
        "The dataset contains 7338 images, each of which is annotated with one label across 5 different flower classes.\n",
        "\n",
        "The input images are stored in a public Cloud Storage bucket. This publicly-accessible bucket also contains a CSV file used to create the Vertex AI multimodal dataset. This file has two columns: the first column lists an image's URI in Cloud Storage, and the second column contains the image's label.\n",
        "\n",
        "In this notebook, we'll use subsets of the flower dataset, each with a fixed number of examples per category, and prepare training, tuning and test subsets\n",
        " as DataFrame.\n",
        "\n",
        "**Tip:** Use the BigFrames library `bpd` instead of `pandas` for larger datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VVl2GGM2RJrU"
      },
      "outputs": [],
      "source": [
        "# Get data from GCS\n",
        "csv = \"gs://cloud-samples-data/ai-platform/flowers/flowers.csv\"\n",
        "all_images = pandas.read_csv(csv, names=[\"image_uris\", \"labels\"])\n",
        "\n",
        "# Shuffle\n",
        "all_images = all_images.sample(frac=1).reset_index(drop=True)\n",
        "\n",
        "# Prepare training and validation set\n",
        "CATEGORIES = [\"daisy\", \"dandelion\", \"roses\", \"sunflowers\", \"tulips\"]\n",
        "# fmt: off\n",
        "TRAINING_CASES_PER_CATEGORY = 100  # @param {type: 'integer'}\n",
        "VALIDATION_CASES_PER_CATEGORY = 100  # @param {type: 'integer'}\n",
        "PREDICTION_DATASET_SIZE = 100  # @param {type: 'integer'}\n",
        "# fmt: on\n",
        "\n",
        "# Set up the prediction set\n",
        "if len(all_images) < PREDICTION_DATASET_SIZE:\n",
        "    raise ValueError(\n",
        "        \"Prediction dataset size is larger than the total number of images.\"\n",
        "    )\n",
        "prediction_set = all_images.iloc[:PREDICTION_DATASET_SIZE]\n",
        "all_images = all_images.iloc[PREDICTION_DATASET_SIZE:]\n",
        "\n",
        "# Set up the training and validation set with evenly distributed labels\n",
        "training_set = pandas.DataFrame()\n",
        "validation_set = pandas.DataFrame()\n",
        "\n",
        "\n",
        "for category in CATEGORIES:\n",
        "    same_labels = all_images[all_images[\"labels\"] == category]\n",
        "    if len(same_labels) < TRAINING_CASES_PER_CATEGORY + VALIDATION_CASES_PER_CATEGORY:\n",
        "        raise ValueError(\"Please reduce the number of cases per category.\")\n",
        "    training_set = pandas.concat(\n",
        "        (training_set, same_labels.iloc[:TRAINING_CASES_PER_CATEGORY]),\n",
        "        ignore_index=True,\n",
        "    )\n",
        "    validation_set = pandas.concat(\n",
        "        (\n",
        "            validation_set,\n",
        "            same_labels.iloc[\n",
        "                TRAINING_CASES_PER_CATEGORY : TRAINING_CASES_PER_CATEGORY\n",
        "                + VALIDATION_CASES_PER_CATEGORY\n",
        "            ],\n",
        "        ),\n",
        "        ignore_index=True,\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MYYGPsA69g_-"
      },
      "outputs": [],
      "source": [
        "# @title Common Functions\n",
        "\n",
        "# Set Pandas display options to show all columns and full width for better inspection\n",
        "pandas.set_option(\"display.max_columns\", None)  # Show all columns\n",
        "pandas.set_option(\"display.expand_frame_repr\", False)  # Prevent line wrapping\n",
        "pandas.set_option(\"display.max_colwidth\", None)  # Show full column width\n",
        "\n",
        "\n",
        "def show_dataset_info(dataset):\n",
        "    \"\"\"Dataset inspection helper\"\"\"\n",
        "    print(\"  Resource name: \", dataset.resource_name)\n",
        "    print(\"  Display name: \", dataset.display_name)\n",
        "    print(\"  Schema URI:   \", dataset.metadata_schema_uri)\n",
        "    print(\"  BQ Table:     \", dataset.bigquery_table)\n",
        "\n",
        "\n",
        "def get_gcs_image(gcs_uri):\n",
        "    \"\"\"Download and show an image from Cloud Storage.\"\"\"\n",
        "    storage_client = storage.Client(project=PROJECT_ID)\n",
        "    blob = storage.blob.Blob.from_string(gcs_uri, client=storage_client)\n",
        "    return Image.open(io.BytesIO(blob.download_as_bytes()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKw7PCHpXwgn"
      },
      "source": [
        "## User Journey Demo\n",
        "\n",
        "The user journey demonstrated here contains the following steps:\n",
        "\n",
        "1. Create Dataset\n",
        "2. Assemble the dataset with a template and inspect assembly\n",
        "3. Run a validation for tuning\n",
        "4. Estimate Resources for tuning\n",
        "5. Run tuning\n",
        "6. Run batch prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HybHbfY85MnR"
      },
      "source": [
        "### 1. Create a dataset from a Pandas or BigFrames DataFrame"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fqewmHHY7TYO"
      },
      "source": [
        "We prepared a DataFrame `training_set` with two columns:\n",
        "\n",
        "*   `image_uris`: GCS URIs of flower images\n",
        "*   `labels`: Flower label (five flower categories, one label per image)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8DxI6amF8ok1"
      },
      "outputs": [],
      "source": [
        "flower_uri = training_set[\"image_uris\"].iloc[0]\n",
        "flower_label = training_set[\"labels\"].iloc[0]\n",
        "\n",
        "display(get_gcs_image(flower_uri))\n",
        "print(f\"Image URI: {flower_uri}\")\n",
        "print(f\"Flower label: {flower_label}\")\n",
        "training_set.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jSxMKdBiId8r"
      },
      "source": [
        "Let's create a Vertex AI multimodal dataset from the prepared DataFrame."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xc37pkL1SnxL"
      },
      "outputs": [],
      "source": [
        "flowers = datasets.MultimodalDataset.from_pandas(dataframe=training_set)\n",
        "\n",
        "# Inspect the dataset\n",
        "show_dataset_info(flowers)\n",
        "flowers.to_bigframes().head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1FfwiGkTgp6F"
      },
      "source": [
        "**Other dataset creation options**\n",
        "\n",
        "Create from a BigQuery table.\n",
        "\n",
        "```py\n",
        "my_dataset_from_bigquery = datasets.MultimodalDataset.from_bigquery(\n",
        "    bigquery_uri=f\"bq://projectId.datasetId.tableId\"\n",
        ")\n",
        "```\n",
        "\n",
        "Create from a BigFrames DataFrame.\n",
        "\n",
        "```py\n",
        "my_dataset_from_pandas = datasets.MultimodalDataset.from_bigframes(\n",
        "    dataframe=my_dataframe\n",
        ")\n",
        "```\n",
        "\n",
        "Create from a GCS file in JSONL format for assembled input (the JSONL file contains Gemini requests, no assembly required).\n",
        "\n",
        "```py\n",
        "my_dataset = datasets.MultimodalDataset.from_gemini_request_jsonl(\n",
        "    gcs_uri=gcs_uri_of_jsonl_file\n",
        ")\n",
        "```\n",
        "\n",
        "List or load existing datasets.\n",
        "\n",
        "```py\n",
        "# Get the most recently created dataset\n",
        "first_dataset = datasets.MultimodalDataset.list()[0]\n",
        "\n",
        "# Load dataset based on dataset name\n",
        "same_dataset = datasets.MultimodalDataset(first_dataset.name)\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vkD3LrKlgkaS"
      },
      "source": [
        "### 2. Assemble the dataset with a template and inspect assembly\n",
        "\n",
        "To use our Flowers dataset with Gemini, let's assemble a full Gemini request referencing the images in our dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PB5VRnpVIEGx"
      },
      "source": [
        "We construct a template configuration by specifying the general prompt, response and system instructions and use placeholders in curly braces. During the assembly, the placeholders are replaced with the values of the dataset column that the placeholders denote.\n",
        "\n",
        "The dataset columns referenced by the placeholders can contain e.g. GCS URIS for files of several data types and modalities:\n",
        "- .pdf\n",
        "- .png, .jpeg, .jpg, .webp\n",
        "- .aac, .flac, .mp3, .m4a, .mpga, .opus, .pcm, .wav\n",
        "- .flv, .mov, .mpegps, .mpg, .wmv, .3pg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gIJUKU6_qV99"
      },
      "outputs": [],
      "source": [
        "template_config = datasets.construct_single_turn_template(\n",
        "    prompt=\"This is the image: {image_uris}\",\n",
        "    response=\"{labels}\",\n",
        "    system_instruction=\"You are a botanical image classifier. Analyze the provided image \"\n",
        "    \"and determine the most accurate classification of the flower.\"\n",
        "    'These are the only flower categories: [/\"daisy/\", /\"dandelion/\", /\"roses/\", /\"sunflowers/\", /\"tulips/\"].'\n",
        "    \"Return only one category per image.\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nD20fEieM4rk"
      },
      "source": [
        "Here, the template is constructed using the library function `construct_single_turn_template()`. Alternatively, it can be explicitly constructed from a Gemini example as below.\n",
        "\n",
        "It is also possible to specify a custom field mapping for the placeholders used in the Gemini example. Then the placeholders can have any name, and not necessarily the column name of the dataset column with the values that are being inserted (here image_uris and labels):\n",
        "\n",
        "```\n",
        "from vertexai.generative_models import Content, Part\n",
        "\n",
        "gemini_example = datasets.GeminiExample(\n",
        "    contents=[\n",
        "        Content(role=\"user\", parts=[Part.from_text(\"This is the image: {uri}\")]),\n",
        "        Content(role=\"model\", parts=[Part.from_text(\"{flower}\")]),\n",
        "    ],\n",
        "    system_instruction=Content(\n",
        "        parts=[\n",
        "            Part.from_text(\n",
        "                'You are a botanical image classifier. Analyze the provided image '\n",
        "                'and determine the most accurate classification of the flower.'\n",
        "                'These are the only flower categories: [/\"daisy/\", /\"dandelion/\", /\"roses/\", /\"sunflowers/\", /\"tulips/\"].'\n",
        "                'Return only one category per image.'\n",
        "            )\n",
        "        ]\n",
        "    ),\n",
        ")\n",
        "\n",
        "template_config = datasets.GeminiTemplateConfig(\n",
        "    gemini_example=gemini_example,\n",
        "    field_mapping={\"uri\": \"image_uris\", \"flower\": \"labels\"},\n",
        ")\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HQnaPYbXPtdI"
      },
      "source": [
        "**Assemble and inspect the dataset.**\n",
        "\n",
        "The dataset assembly creates a BigQuery table with the assembled examples in a single `request` column. The assembly method below returns a tuple containing a table id (`str`) referencing the assembly BQ table, and a DataFrame (`bigframes.pandas.DataFrame`) for direct inspection.\n",
        "The DataFrame and the BQ table referenced by the table id contain the assembled dataset in a single column `request`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A5zGxYOYQSJE"
      },
      "outputs": [],
      "source": [
        "table_id, assembly = flowers.assemble(template_config=template_config)\n",
        "\n",
        "# Inspect assembled dataset\n",
        "assembly.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S0KeOWRvSj75"
      },
      "source": [
        "It is also possible to attach the template and run the assembly and the validation below without passing it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VE2txCcS6Hlf"
      },
      "outputs": [],
      "source": [
        "flowers.attach_template_config(template_config=template_config)\n",
        "_, _ = flowers.assemble()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lk8zIbaQ55Bp"
      },
      "source": [
        "### 3. Run a validation for tuning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m7UODTevZovT"
      },
      "source": [
        "Validate a dataset for tuning.\n",
        "Tuning dataset usages are: `SFT_VALIDATION`, `SFT_TRAINING`.\n",
        "\n",
        "First we attach the `template_config` and use it implicitly for all further tasks."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_4itps2hkqVS"
      },
      "outputs": [],
      "source": [
        "validation = flowers.assess_tuning_validity(\n",
        "    model_name=\"gemini-2.0-flash-001\", dataset_usage=\"SFT_TRAINING\"\n",
        ")\n",
        "\n",
        "# Check if there are validation errors\n",
        "validation.errors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tt1n6Mwak864"
      },
      "source": [
        "Let's validate a dataset with an incorrect `template_config`, e.g. using a `GeminiExample` that contains two consecutive `user` contents, instead of a `user` content followed by a `model` content."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yscl1KH8XUTy"
      },
      "outputs": [],
      "source": [
        "invalid_gemini_example = datasets.GeminiExample(\n",
        "    contents=[\n",
        "        Content(role=\"user\", parts=[Part.from_text(\"This is the image: {image_uris}\")]),\n",
        "        # Consecutive content turn with the same role\n",
        "        Content(role=\"user\", parts=[Part.from_text(\".\")]),\n",
        "    ],\n",
        ")\n",
        "invalid_configuration = datasets.GeminiTemplateConfig(\n",
        "    gemini_example=invalid_gemini_example\n",
        ")\n",
        "\n",
        "validation = flowers.assess_tuning_validity(\n",
        "    model_name=\"gemini-2.0-flash-001\",\n",
        "    dataset_usage=\"SFT_TRAINING\",\n",
        "    template_config=invalid_configuration,\n",
        ")\n",
        "\n",
        "validation.errors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ayqG3aKX55P8"
      },
      "source": [
        "### 4. Estimate resources for tuning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s8ufb8KWYE51"
      },
      "outputs": [],
      "source": [
        "tuning_resources = flowers.assess_tuning_resources(model_name=\"gemini-2.5-flash-001\")\n",
        "print(tuning_resources)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i5wAETJ055XQ"
      },
      "source": [
        "## 5. Run Tuning\n",
        "\n",
        "Prerequisites:\n",
        "\n",
        "- Your Vertex service account `service-{project_number}@gcp-sa-vertex-tune.iam.gserviceaccount.com` needs to have read permissions on to the GCS buckets referenced in the dataset. If this is not automatically the case then you might have to assign the service account a role such as `Storage Object User`, see the screenshot below.\n",
        "\n",
        "- The Vertex multimodal dataset needs to have an attached `template_config` (run `flowers.attach_template_config(template_config=template_config)`)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H9E_MPmO6hLi"
      },
      "source": [
        "Let's also prepare and use the validation dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XIzeW5Qw36Si"
      },
      "outputs": [],
      "source": [
        "# Optional: Create Vertex Multimodal dataset for the validation set\n",
        "flowers_validation = datasets.MultimodalDataset.from_pandas(dataframe=validation_set)\n",
        "\n",
        "# Attach the template config\n",
        "flowers_validation.attach_template_config(template_config=template_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7WdVArD16oHZ"
      },
      "source": [
        "Here we use the training and validation set to start a tuning job:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M62UKX-3byET"
      },
      "outputs": [],
      "source": [
        "from vertexai.preview.tuning import sft\n",
        "\n",
        "tuning_job = sft.train(\n",
        "    source_model=\"gemini-2.0-flash-001\",\n",
        "    train_dataset=flowers,\n",
        "    validation_dataset=flowers_validation,  # optional\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y41fcy6EII5C"
      },
      "source": [
        "Let's monitor the job state and obtain the model ID of the tuned model once the tuning job has ended."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rZiHdiB4IGyS"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "print(f\"Tuning job started: {tuning_job.resource_name}\")\n",
        "\n",
        "# Wait for the job to complete\n",
        "while tuning_job.has_ended is False:\n",
        "    time.sleep(60)\n",
        "    tuning_job.refresh()  # Refresh the job state\n",
        "    print(f\"Polling - Current job state: {tuning_job.state}\")\n",
        "\n",
        "# Check the final state\n",
        "if tuning_job.state == \"JOB_STATE_FAILED\":\n",
        "    print(f\"Tuning job failed: {tuning_job.error}\")\n",
        "else:\n",
        "    print(f\"Tuning job ended with state: {tuning_job.state}\")\n",
        "    # Get model ID\n",
        "    tuned_model_id = tuning_job.tuned_model_name.split(\"/\")[-1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2HVR2huE609v"
      },
      "source": [
        "The tuning job can also be started from the GenAI SDK:\n",
        "\n",
        "\n",
        "```\n",
        "from google import genai\n",
        "from google.genai.types import HttpOptions, CreateTuningJobConfig\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "tuning_job = client.tunings.tune(\n",
        "    base_model=\"gemini-2.5-flash\",\n",
        "    training_dataset = {\n",
        "        \"vertex_dataset_resource\":flowers.resource_name\n",
        "    },\n",
        "    config=CreateTuningJobConfig(\n",
        "        tuned_model_display_name=\"Example tuning job with Multimodal Dataset 'Flowers'\",\n",
        "        validation_dataset = {\n",
        "        \"vertex_dataset_resource\":flowers_validation.resource_name\n",
        "    },\n",
        "    ),\n",
        ")\n",
        "\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Rw4BE0t7Jzd"
      },
      "source": [
        "## 6. Batch Prediction\n",
        "\n",
        "You can use Vertex Multimodal datasets to run a batch prediction job with Gemini.\n",
        "In the future we will support directly passing Vertex Multimodal dataset objects, for now users can simply pass the BigQuery URI of the assembly:\n",
        "\n",
        "```py\n",
        "# Currently\n",
        "\n",
        "prediction_assembly_table_id, _ = flowers_prediction.assemble(template_config=template_config)\n",
        "\n",
        "job = client.batches.create(\n",
        "    model=model,\n",
        "    src=f\"bq://{prediction_assembly_table_id}\",\n",
        ")\n",
        "\n",
        "# After the integration\n",
        "\n",
        "job = client.batches.create(\n",
        "    model=model,\n",
        "    src=flowers_prediction,\n",
        ")\n",
        "\n",
        "```\n",
        "\n",
        "See also: [Batch Prediction Documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini) - Note that Vertex Multimodal datasets will be added as another dataset source next to Google Cloud Storage and BigQuery in the documentation soon, when the integration has been completed.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "skq_VSZT9-sj"
      },
      "source": [
        "The model specified can either be a Gemini base model, the model you tuned above, or any other tuned custom model.\n",
        "\n",
        "By default the colab uses the model tuned above for the batch prediction. If you want to run a batch prediction job on a base model or on another custom model, you can provide the base model name or the custom model id, respectively, in the field below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5K1dGtdg97j5"
      },
      "outputs": [],
      "source": [
        "# @title Specify the Model used for Batch Prediction\n",
        "# fmt: off\n",
        "model_for_batch_prediction = \"tuned model\"  # @param [\"tuned model\", \"base model\", \"other custom model\"]\n",
        "# fmt: on\n",
        "optional_base_model_name = \"\"  # @param {type:\"string\"}\n",
        "optional_custom_model_id = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set full model name\n",
        "if model_for_batch_prediction == \"tuned model\":\n",
        "    model = f\"projects/{PROJECT_ID}/locations/{LOCATION}/models/{tuned_model_id}\"\n",
        "elif model_for_batch_prediction == \"base model\":\n",
        "    if not optional_base_model_name:\n",
        "        raise ValueError(\n",
        "            \"Please provide a optional_base_model_name when 'base model' is selected.\"\n",
        "        )\n",
        "    model = optional_base_model_name\n",
        "elif model_for_batch_prediction == \"other custom model\":\n",
        "    if not optional_custom_model_id:\n",
        "        raise ValueError(\n",
        "            \"Please provide a optional_custom_model_id when 'other custom model' is selected.\"\n",
        "        )\n",
        "    model = (\n",
        "        f\"projects/{PROJECT_ID}/locations/{LOCATION}/models/{optional_custom_model_id}\"\n",
        "    )\n",
        "\n",
        "\n",
        "print(f\"Using model: {model}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmvL3rHCB9c4"
      },
      "source": [
        "Let's prepare the prediction set as Vertex Multimodal dataset and start a Batch Prediction job with the specified model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dufR8xxEqorc"
      },
      "outputs": [],
      "source": [
        "# Get Vertex Multimodal Dataset\n",
        "flowers_prediction = datasets.MultimodalDataset.from_pandas(dataframe=prediction_set)\n",
        "\n",
        "# Assemble the dataset with a template config and get the assembly table id\n",
        "prediction_assembly_table_id, _ = flowers_prediction.assemble(\n",
        "    template_config=template_config\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qr1Xtk5Przaa"
      },
      "outputs": [],
      "source": [
        "# @title Run a Batch Prediction Job\n",
        "\n",
        "from google import genai\n",
        "from google.genai.types import JobState\n",
        "\n",
        "# Initialize GenAI client for batch prediction library\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "job = client.batches.create(\n",
        "    # use the model selected above\n",
        "    model=model,\n",
        "    # Use the assembly table id as source\n",
        "    src=f\"bq://{prediction_assembly_table_id}\",\n",
        ")\n",
        "\n",
        "\n",
        "completed_states = {\n",
        "    JobState.JOB_STATE_SUCCEEDED,\n",
        "    JobState.JOB_STATE_FAILED,\n",
        "    JobState.JOB_STATE_CANCELLED,\n",
        "    JobState.JOB_STATE_PAUSED,\n",
        "}\n",
        "\n",
        "while job.state not in completed_states:\n",
        "    time.sleep(30)\n",
        "    job = client.batches.get(name=job.name)\n",
        "    print(f\"Job state: {job.state}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "intro_vertex_ai_multimodal_dataset.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
