notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "OdZIyZwjgsQcOXnmE8X0xy40",
      "metadata": {
        "cellView": "form",
        "id": "OdZIyZwjgsQcOXnmE8X0xy40"
      },
      "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",
      "id": "VJWDivOv3OWy",
      "language": "markdown",
      "metadata": {
        "id": "VJWDivOv3OWy"
      },
      "source": [
        "# Vertex AI Model Garden - PaliGemma (Finetuning)\n",
        "\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_jax_paligemma_finetuning.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eMKfcPJVXlWR"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to do finetuning PaliGemma with a Vertex AI Custom Training Job, deploying the finetuned model to a Vertex AI Endpoint, and making online predictions.\n",
        "\n",
        "\n",
        "### Objective\n",
        "- Prepare data for finetuning.\n",
        "- Launch a Vertex AI Custom Training Job to finetune PaliGemma, storing the resulting model to a GCS bucket.\n",
        "- Deploy the finetuned PaliGemma model to a Vertex AI Endpoint.\n",
        "- Make predictions to the endpoint.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2aFHbs1g6Wc-",
      "metadata": {
        "id": "2aFHbs1g6Wc-"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "QvQjsmIJ6Y3f",
      "metadata": {
        "cellView": "form",
        "id": "QvQjsmIJ6Y3f"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
        "\n",
        "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
        "\n",
        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
        "# @markdown | ----------- | ----------- | ----------- |\n",
        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
        "\n",
        "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import json\n",
        "import os\n",
        "import tempfile\n",
        "import uuid\n",
        "from typing import Tuple\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
        "models, endpoints = {}, {}\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "if not REGION:\n",
        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
        "        raise ValueError(\n",
        "            \"REGION must be set. See\"\n",
        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
        "            \" available cloud locations.\"\n",
        "        )\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"paligemma\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
        "\n",
        "pretrained_filename_lookup = {\n",
        "    \"paligemma-224-float32\": \"pt_224.npz\",\n",
        "    \"paligemma-448-float32\": \"pt_448.npz\",\n",
        "    \"paligemma-896-float32\": \"pt_896.npz\",\n",
        "    \"paligemma-mix-224-float32\": \"mix_224.npz\",\n",
        "    \"paligemma-mix-448-float32\": \"mix_448.npz\",\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "Jo9Iv7WZGQub",
      "metadata": {
        "cellView": "form",
        "id": "Jo9Iv7WZGQub"
      },
      "outputs": [],
      "source": [
        "# @title Access PaliGemma models on Vertex AI for GPU based serving\n",
        "\n",
        "\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [PaliGemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/363) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
        "# @markdown 1. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
        "# @markdown 1. After accepting the agreement of PaliGemma, a `gs://` URI containing PaliGemma pretrained models will be shared.\n",
        "# @markdown 1. Paste the link in the `VERTEX_AI_MODEL_GARDEN_PALIGEMMA` field below.\n",
        "# @markdown 1. The PaliGemma models will be copied into `BUCKET_URI`.\n",
        "# @markdown The file transfer can take anywhere from 15 minutes to 30 minutes.\n",
        "VERTEX_AI_MODEL_GARDEN_PALIGEMMA = \"gs://\"  # @param {type:\"string\", isTemplate:true}\n",
        "assert (\n",
        "    VERTEX_AI_MODEL_GARDEN_PALIGEMMA and VERTEX_AI_MODEL_GARDEN_PALIGEMMA != \"gs://\"\n",
        "), \"Click the agreement of PaliGemma in Vertex AI Model Garden, and get the GCS path of PaliGemma model artifacts.\"\n",
        "print(\n",
        "    \"Copying PaliGemma model artifacts from\",\n",
        "    VERTEX_AI_MODEL_GARDEN_PALIGEMMA,\n",
        "    \"to \",\n",
        "    MODEL_BUCKET,\n",
        ")\n",
        "\n",
        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_PALIGEMMA/* $MODEL_BUCKET\n",
        "\n",
        "assert (\n",
        "    os.system(f\"gsutil ls {MODEL_BUCKET}\") == 0\n",
        "), f\"MODEL_BUCKET does not exist: {MODEL_BUCKET}.\"\n",
        "model_path_prefix = MODEL_BUCKET"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "XdVYFARuNsuG",
      "metadata": {
        "id": "XdVYFARuNsuG"
      },
      "source": [
        "## Finetune with Vertex AI Custom Training Jobs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "zSjxLZegONqH",
      "metadata": {
        "cellView": "form",
        "id": "zSjxLZegONqH"
      },
      "outputs": [],
      "source": [
        "# @title Data preparation\n",
        "\n",
        "# @markdown The dataset file format is to be image-string pairs stored in a jsonl file.\n",
        "# @markdown The value for `\"image\"` can be a GCS path or a URL.\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown {\"image\": \"gs://bucket-name/image.jpg\", \"prefix\": \"What animal is this?\", \"suffix\": \"cat\"}\n",
        "# @markdown {\"image\": \"https://google.com/image.jpg\", \"prefix\": \"What drink is this?\", \"suffix\": \"soda\"}\n",
        "# @markdown ```\n",
        "\n",
        "dataset_gcs_uri = \"gs://longcap100/data_train90.jsonl\"  # @param {type: \"string\"}\n",
        "\n",
        "# @markdown [Optional] You can optionally specify the image fields in the JSONL file to use the\n",
        "# @markdown filename and fill in the `dataset_image_dir` with the location where the images are stored.\n",
        "dataset_image_dir = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YMteCx6OS_bq"
      },
      "source": [
        "### Inference examples before finetuning\n",
        "\n",
        "The images below are some of the examples of inference results of the pretrained\n",
        " `paligemma-224-float32` checkpoint.\n",
        "\n",
        "| Image | URI | Caption |\n",
        "|-----|-----|-----|\n",
        "| <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"200\" >  | gs://longcap100/91.jpeg | the beauty of the sleeve |\n",
        "| <img src=\"https://storage.googleapis.com/longcap100/92.jpeg\" width=\"200\" >   | gs://longcap100/92.jpeg | how to wear a maxi dress for summer |\n",
        "| <img src=\"https://storage.googleapis.com/longcap100/93.jpeg\" width=\"200\" >  | gs://longcap100/93.jpeg | a red blazer and black bag , a key piece of the week 's fashion . |\n",
        "| <img src=\"https://storage.googleapis.com/longcap100/94.jpeg\" width=\"200\" >   | gs://longcap100/94.jpeg | how to wear boyfriend jeans like a fashion blogger |\n",
        "| <img src=\"https://storage.googleapis.com/longcap100/95.jpeg\" width=\"200\" >   | gs://longcap100/95.jpeg | this graphic sweatshirt is a must have for your wardrobe . |\n",
        "| <img src=\"https://storage.googleapis.com/longcap100/96.jpeg\" width=\"200\" >   | gs://longcap100/96.jpeg | person in a long shot of our model |\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k_ET-ZVpWFVk"
      },
      "source": [
        "### Finetune"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "jitdsEHDNqQk",
      "metadata": {
        "cellView": "form",
        "id": "jitdsEHDNqQk"
      },
      "outputs": [],
      "source": [
        "# @title Run\n",
        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs. With the default setting, you can train about 30 steps per minute with the batch size of 4.\n",
        "\n",
        "model_variant = \"pt\"  # @param [\"mix\", \"pt\"]\n",
        "model_resolution = 224  # @param [224, 448, 896]\n",
        "model_precision_type = \"float32\"  # Only float32 is supported.\n",
        "\n",
        "if model_variant == \"mix\":\n",
        "    model_name_prefix = \"paligemma-mix\"\n",
        "else:\n",
        "    model_name_prefix = \"paligemma\"\n",
        "\n",
        "base_model_name = f\"{model_name_prefix}-{model_resolution}-{model_precision_type}\"\n",
        "base_model_filename = pretrained_filename_lookup[base_model_name]\n",
        "base_model_uri = os.path.join(model_path_prefix, base_model_filename)\n",
        "\n",
        "# The pre-built training docker image.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-paligemma-train-gpu:20240807_0916_RC00\"\n",
        "\n",
        "# The accelerator to use.\n",
        "ACCELERATOR_TYPE = \"NVIDIA_L4\"  # @param [\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\"]\n",
        "\n",
        "# Batch size for finetuning.\n",
        "batch_size = 4  # @param {type:\"integer\"}\n",
        "# Number of epochs to train.\n",
        "epochs = 3  # @param {type:\"integer\"}\n",
        "# Learning rate.\n",
        "learning_rate = 0.1  # @param{type:\"number\"}\n",
        "# Text length.\n",
        "text_length = 512  # @param{type:\"integer\"}\n",
        "\n",
        "# Worker pool spec.\n",
        "\n",
        "if ACCELERATOR_TYPE == \"NVIDIA_TESLA_V100\":\n",
        "    machine_type = \"n1-standard-8\"\n",
        "    accelerator_count = 2\n",
        "elif ACCELERATOR_TYPE == \"NVIDIA_L4\":\n",
        "    machine_type = \"g2-standard-24\"\n",
        "    accelerator_count = 2\n",
        "else:\n",
        "    raise ValueError(\n",
        "        f\"Cannot automatically determine machine type from {ACCELERATOR_TYPE}.\"\n",
        "    )\n",
        "\n",
        "replica_count = 1\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=True,\n",
        ")\n",
        "\n",
        "# Setup training job.\n",
        "job_name = common_util.get_job_name_with_datetime(\"paligemma-finetune\")\n",
        "\n",
        "# Add labels for the finetuning job.\n",
        "labels = {\n",
        "    \"mg-source\": \"notebook\",\n",
        "    \"mg-notebook-name\": \"model_garden_jax_paligemma_finetuning.ipynb\".split(\".\")[0],\n",
        "}\n",
        "\n",
        "labels[\"mg-tune\"] = \"publishers-google-models-paligemma\"\n",
        "versioned_model_id = base_model_name.lower().replace(\".\", \"-\")\n",
        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        "    labels=labels,\n",
        ")\n",
        "\n",
        "# Designate a GCS folder to store the LORA adapter.\n",
        "finetune_output_dir_name = common_util.get_job_name_with_datetime(\"paligemma-finetune\")\n",
        "finetune_output_dir = os.path.join(STAGING_BUCKET, finetune_output_dir_name)\n",
        "\n",
        "train_args = [\n",
        "    \"--config=big_vision/configs/proj/paligemma/transfers/vertexai_l4.py\",\n",
        "    f\"--workdir={finetune_output_dir}\",\n",
        "    f\"--config.model_init={base_model_uri}\",\n",
        "    f\"--config.input.data.fname={dataset_gcs_uri}\",\n",
        "    f\"--config.text_len={text_length}\",\n",
        "]\n",
        "\n",
        "if batch_size:\n",
        "    train_args.append(f\"--config.input.batch_size={batch_size}\")\n",
        "if epochs:\n",
        "    train_args.append(f\"--config.total_epochs={epochs}\")\n",
        "if learning_rate:\n",
        "    train_args.append(f\"--config.lr={learning_rate}\")\n",
        "\n",
        "if dataset_image_dir:\n",
        "    train_args.append(f\"--config.input.data.fopen_keys.image={dataset_image_dir}\")\n",
        "train_job.run(\n",
        "    args=train_args,\n",
        "    replica_count=replica_count,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=accelerator_count,\n",
        "    boot_disk_size_gb=500,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")\n",
        "\n",
        "print(\"Checkpoint and log files was saved in: \", finetune_output_dir)\n",
        "\n",
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "WngKdCIzOTJV",
      "metadata": {
        "cellView": "form",
        "id": "WngKdCIzOTJV"
      },
      "outputs": [],
      "source": [
        "# @title View training loss\n",
        "# @markdown Metrics will be stored in a file named `big_vision_metrics.txt` in the GCS bucket, including training loss, examples seen, and core hours throughout training.\n",
        "\n",
        "# @markdown Run this cell to get and plot the training loss.\n",
        "\n",
        "# Get relevant metrics from metrics file.\n",
        "metrics_file_name = \"big_vision_metrics.txt\"\n",
        "metrics_path = os.path.join(finetune_output_dir, metrics_file_name)\n",
        "\n",
        "temp_dir = tempfile.TemporaryDirectory()\n",
        "local_metrics_path = os.path.join(temp_dir.name, metrics_file_name)\n",
        "\n",
        "! gsutil cp $metrics_path $local_metrics_path\n",
        "\n",
        "steps = []\n",
        "training_losses = []\n",
        "with open(local_metrics_path, \"r\") as f:\n",
        "    for line in f:\n",
        "        metric = json.loads(line)\n",
        "        steps.append(metric[\"step\"])\n",
        "        training_losses.append(metric[\"training_loss\"])\n",
        "\n",
        "# Plot training plot\n",
        "plt.plot(steps, training_losses)\n",
        "plt.title(\"Steps vs. Training Loss\")\n",
        "plt.xlabel(\"Steps\")\n",
        "plt.ylabel(\"Training Loss\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uVKr3rggXJd0"
      },
      "source": [
        "## Deployment and prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "toY-WPKDFesF"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "# @markdown This section uploads the finetuned PaliGemma model to Model Registry and deploys it to a Vertex AI Endpoint. It takes approximately 15 minutes to finish.\n",
        "\n",
        "# @markdown Note: You cannot use accelerator type `NVIDIA_TESLA_V100` to serve prebuilt or finetuned PaliGemma models with resolution `896`.\n",
        "\n",
        "last_checkpoint_file_name = \"checkpoint.bv-LAST\"\n",
        "last_checkpoint_path = os.path.join(finetune_output_dir, last_checkpoint_file_name)\n",
        "\n",
        "local_last_checkpoint_path = os.path.join(temp_dir.name, last_checkpoint_file_name)\n",
        "\n",
        "! gsutil cp $last_checkpoint_path $local_last_checkpoint_path\n",
        "\n",
        "with open(local_last_checkpoint_path, \"r\") as f:\n",
        "    final_checkpoint_name = \"checkpoint.bv-\" + f.read()\n",
        "    checkpoint_path = os.path.join(finetune_output_dir, final_checkpoint_name)\n",
        "\n",
        "model_name = f\"paligemma-{model_resolution}-{model_precision_type}-custom\"\n",
        "print(f\"Deploying custom PaliGemma model: {model_name}\")\n",
        "\n",
        "# The pre-built serving docker images.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-paligemma-serve-gpu:20240807_0916_RC00\"\n",
        "\n",
        "# @markdown Select the accelerator type to use to deploy the model:\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
        "if accelerator_type == \"NVIDIA_L4\":\n",
        "    machine_type = \"g2-standard-16\"\n",
        "    accelerator_count = 1\n",
        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
        "    if model_resolution == 896 and model_precision_type == \"float32\":\n",
        "        raise ValueError(\n",
        "            \"NVIDIA_TESLA_V100 is not sufficient. Multi-gpu is not supported for PaLIGemma.\"\n",
        "        )\n",
        "    else:\n",
        "        machine_type = \"n1-highmem-8\"\n",
        "        accelerator_count = 1\n",
        "else:\n",
        "    raise ValueError(\n",
        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
        "    )\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "\n",
        "def deploy_model(\n",
        "    model_name: str,\n",
        "    checkpoint_path: str,\n",
        "    machine_type: str = \"g2-standard-32\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    resolution: int = 224,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
        "    model_name_with_time = common_util.get_job_name_with_datetime(model_name)\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name_with_time}-endpoint\"\n",
        "    )\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name_with_time,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/predict\",\n",
        "        serving_container_health_route=\"/health\",\n",
        "        serving_container_environment_variables={\n",
        "            \"CKPT_PATH\": checkpoint_path,\n",
        "            \"RESOLUTION\": resolution,\n",
        "            \"MODEL_ID\": model_name,\n",
        "        },\n",
        "        model_garden_source_model_name=\"publishers/google/models/paligemma\",\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name_with_time} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        enable_access_logging=True,\n",
        "        min_replica_count=1,\n",
        "        sync=True,\n",
        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_jax_paligemma_finetuning.ipynb\"},\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
        "models[\"model\"], endpoints[\"endpoint\"] = deploy_model(\n",
        "    model_name=model_name,\n",
        "    checkpoint_path=checkpoint_path,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    resolution=model_resolution,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G_FT3hglXMks"
      },
      "source": [
        "### Image captioning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0v3GA8NU8SNp",
      "metadata": {
        "cellView": "form",
        "id": "0v3GA8NU8SNp"
      },
      "outputs": [],
      "source": [
        "# @markdown This section uses the deployed PaliGemma model to caption and describe an image in a chosen language. Check how the caption has changed compared to the examples above.\n",
        "\n",
        "caption_prompt = True\n",
        "\n",
        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"400\" >\n",
        "\n",
        "image_url = \"https://storage.googleapis.com/longcap100/91.jpeg\"  # @param {type:\"string\"}\n",
        "language_code = \"en\"  # @param {type: \"string\"}\n",
        "\n",
        "image = common_util.download_image(image_url)\n",
        "display(image)\n",
        "\n",
        "# Make a prediction.\n",
        "image_base64 = common_util.image_to_base64(image)\n",
        "\n",
        "caption = common_util.caption_predict(\n",
        "    endpoints[\"endpoint\"], language_code, image, caption_prompt\n",
        ")\n",
        "\n",
        "print(\"Caption: \", caption)\n",
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "IrVZ030i4XMY",
      "metadata": {
        "id": "IrVZ030i4XMY"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YsMpOI1kYjil"
      },
      "outputs": [],
      "source": [
        "# Delete the train job.\n",
        "train_job.delete()\n",
        "\n",
        "# Clean up the temporary directory\n",
        "temp_dir.cleanup()\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "for endpoint in endpoints.values():\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "for model in models.values():\n",
        "    model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_jax_paligemma_finetuning.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}