notebooks/community/gapic/custom/showcase_custom_image_classification_online_explain_example_based_api.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": "pS9oVc0bScpN"
      },
      "source": [
        "# Vertex client library: Custom training image classification model for online prediction with explanation using Example-based API \n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/gapic/custom/showcase_custom_image_classification_online_explain_example_based_api.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/gapic/custom/showcase_custom_image_classification_online_explain_example_based_api.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/notebooks/community/gapic/custom/showcase_custom_image_classification_online_explain_example_based_api.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to get example-based explanations for your model. Explanations can help you answer questions about why a model made a certain prediction, relating it to the features in the training data. In this demo we will go over:\n",
        "\n",
        "1. Getting Example-Based explanations from Vertex Explainable AI services.\n",
        "2. A use-case for exploring similar examples to understand model predictions.\n",
        "\n",
        "The prerequisites for this notebook are:\n",
        "1. A predictive model and a way to extract a latent representation from it -- i.e., embedding. The notebook will demonstrate how this can be done for a Deep Neural Network.\n",
        "2. A Google Cloud project.\n",
        "3. A Google bucket to host the model and the dataset.\n",
        "\n",
        "Once these are in place, the three main sections of the notebook are:\n",
        "1. Creating and uploading a model with explanations enabled.\n",
        "2. Creating an `Endpoint` resource and deploying the model to it.\n",
        "3. Issuing explanation request and inspecting them.\n",
        "\n",
        "### Dataset\n",
        "\n",
        "For this notebook, we will use the [STL10 dataset](https://cs.stanford.edu/~acoates/stl10/) downloaded through [TF Datasets](https://www.tensorflow.org/datasets/catalog/stl10). This dataset is a subset of the [ImageNet dataset](https://www.image-net.org/) with just 10 classes.\n",
        "\n",
        "### Objective\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI Training`\n",
        "- `Vertex Explainable AI`\n",
        "- `Vertex AI Prediction`\n",
        "- `Vertex AI Model`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "1. Prepare training data\n",
        "2. Fine tune a image classication model to get embeddings\n",
        "3. Register the model in Vertex AI Model Registry\n",
        "4. Deploy the model in Vertex AI Endpoint\n",
        "5. Request explanations using Example-Based Explanation API\n",
        "6. Analyze the results\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\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ze4-nDLfK4pw"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or Google Cloud Notebooks, you can skip this step**, since your environment already meets\n",
        "all the requirements to run this notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCuSR8GkAgzl"
      },
      "source": [
        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
        "You need the following:\n",
        "\n",
        "* The Google Cloud SDK\n",
        "* Git\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps provide a condensed set of\n",
        "instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip3 install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "## Installations\n",
        "\n",
        "Install the following packages to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Vertex AI Workbench Notebook product has specific requirements\n",
        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
        "    \"/opt/deeplearning/metadata/env_version\"\n",
        ")\n",
        "\n",
        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_WORKBENCH_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\"\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! pip3 install {USER_FLAG} --upgrade numpy tensorflow tensorflow_datasets -q\n",
        "    ! pip3 install {USER_FLAG} --upgrade google-cloud-aiplatform -q\n",
        "else:\n",
        "    ! pip3 install {USER_FLAG} --upgrade numpy==1.21.6 tensorflow==2.8.0 tensorflow_datasets==4.6.0 -q\n",
        "    ! pip3 install {USER_FLAG} --upgrade google-cloud-aiplatform==1.15.0 -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhq5zEbGg0XX"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install/upgrade the packages, you need to restart the notebook kernel so it can find those packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lWEdiXsJg0XY"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,cloudresourcemanager.googleapis.com)\n",
        "\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7QEa15y1hofD"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"\"\n",
        "\n",
        "# Get your Google Cloud project ID from gcloud\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID: \", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eTnyUHeOKb-h"
      },
      "source": [
        "Otherwise, set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_KnmhsDAKc5q"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "! gcloud config set project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Aepw5iW5Ko0s"
      },
      "source": [
        "#### Get your project number \n",
        "\n",
        "Now that the project ID is set, you get your corresponding project number."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OXcDgJhlKw9R"
      },
      "outputs": [],
      "source": [
        "shell_output = ! gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
        "PROJECT_NUMBER = shell_output[0]\n",
        "print(\"Project Number:\", PROJECT_NUMBER)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "### Region\n",
        "\n",
        "You can also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
        "\n",
        "- Americas: `us-central1`\n",
        "- Europe: `europe-west4`\n",
        "- Asia Pacific: `asia-east1`\n",
        "\n",
        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
        "\n",
        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sKBTnvJpox9P"
      },
      "outputs": [],
      "source": [
        "REGION = \"[your-region]\"  # @param {type:\"string\"}\n",
        "\n",
        "if REGION == \"[your-region]\":\n",
        "    REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "### Timestamp\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "697568e92bd6"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dr--iN2kAylZ"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using Vertex AI Workbench Notebooks, you will skip this step**, since your environment is already\n",
        "authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "**If you are using Colab**, run the cell below and follow the instructions\n",
        "when prompted to authenticate your account via oAuth.\n",
        "\n",
        "**Otherwise**, follow these steps:\n",
        "\n",
        "1. In the Cloud Console, go to the [**Create service account key**\n",
        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "\n",
        "2. Click **Create service account**.\n",
        "\n",
        "3. In the **Service account name** field, enter a name, and\n",
        "   click **Create**.\n",
        "\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type the following role into the filter box and select them\n",
        "\n",
        "    - Service Account User\n",
        "    - Storage Admin\n",
        "    - Storage Object Admin\n",
        "    - Vertex AI Administrator\n",
        "\n",
        "\n",
        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
        "local environment.\n",
        "\n",
        "6. Enter the path to your service account key as the\n",
        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "# If you are running this notebook in Colab, run this cell and follow the\n",
        "# instructions to authenticate your GCP account. This provides access to your\n",
        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
        "# requests.\n",
        "\n",
        "import sys\n",
        "\n",
        "# If on Vertex AI Workbench, then don't execute this code\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
        "    \"DL_ANACONDA_HOME\"\n",
        "):\n",
        "    if \"google.colab\" in sys.modules:\n",
        "        from google.colab import auth as google_auth\n",
        "\n",
        "        google_auth.authenticate_user()\n",
        "\n",
        "    # If you are running this notebook locally, replace the string below with the\n",
        "    # path to your service account key and run this cell to authenticate your GCP\n",
        "    # account.\n",
        "    elif not os.getenv(\"IS_TESTING\"):\n",
        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
        "Cloud Storage buckets.\n",
        "\n",
        "You may also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook. We suggest that you [choose a region where Vertex AI services are\n",
        "available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bHb-DP1ieozQ"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf221059d072"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID + \"-aip-\" + TIMESTAMP\n",
        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Run the following cell to create your Cloud Storage bucket.** If the bucket already exists, you will get an error but it wouldn't affect the rest of the tutorial. However, you might get unwanted data in this bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ucvCsknMCims"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents. If this is a new bucket, this cell wouldn't produce an output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhOb7YnwClBb"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "source": [
        "### Get your service account \n",
        "\n",
        "If you do not want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vqUH5pk0OWzl"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "autoset_service_account"
      },
      "outputs": [],
      "source": [
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if not IS_COLAB:\n",
        "        shell_output = !gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    else:  # IS_COLAB:\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",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account:pipelines"
      },
      "source": [
        "### Set service account access \n",
        "\n",
        "Run the following commands to grant your service account access to the bucket that you created in the previous step. You only need to run this step once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nGfP_UQEOWzm"
      },
      "outputs": [],
      "source": [
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator {BUCKET_URI}\n",
        "\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sYb3wtPucIUR"
      },
      "source": [
        "### Create local directories\n",
        "Next, you create some local directories that you use in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iHd3PHCFcHKj"
      },
      "outputs": [],
      "source": [
        "DATA_PATH = \"data\"\n",
        "DELIVERABLE_PATH = \"deliverables\"\n",
        "\n",
        "! mkdir -m 777 -p {DATA_PATH}\n",
        "! mkdir -m 777 -p {DELIVERABLE_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y9Uo3tifg1kx"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import io\n",
        "import json\n",
        "# General\n",
        "import time\n",
        "\n",
        "# Training\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "%matplotlib inline\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "# Vertex AI\n",
        "from google.cloud import aiplatform_v1beta1 as vertex_ai_v1beta1\n",
        "from google.cloud.aiplatform_v1beta1.types import io as io_pb2\n",
        "from google.protobuf import json_format\n",
        "from google.protobuf.struct_pb2 import Value\n",
        "from PIL import Image\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Set up variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ikn2WUNWPG2B"
      },
      "source": [
        "#### Set Tutorial variables\n",
        "\n",
        "Set some variables for the tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-P3O55RyMGtC"
      },
      "outputs": [],
      "source": [
        "# General\n",
        "ENVIRON = \"prod\"\n",
        "DATASET_NAME = (\n",
        "    \"stl10\"  # Will be downloaded from https://www.tensorflow.org/datasets/catalog/stl10\n",
        ")\n",
        "\n",
        "# Model experimentation\n",
        "RAW_DIR = f\"{DATA_PATH}/raw/{DATASET_NAME}\"\n",
        "PREPROCESSED_DIR = f\"{DATA_PATH}/preprocessed/{DATASET_NAME}\"\n",
        "MODEL_DIR = f\"{DELIVERABLE_PATH}/models/mobilenetv2-{DATASET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NGFeDbldOr6y"
      },
      "source": [
        "#### Set Vertex AI constants\n",
        "\n",
        "Setup up the following constants for Vertex:\n",
        "\n",
        "- `API_ENDPOINT`: The Vertex API service endpoint for dataset, model, job, pipeline and endpoint services.\n",
        "- `PARENT`: The Vertex location root path for dataset, model, job, pipeline and endpoint resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HN-kbHKvOw7G"
      },
      "outputs": [],
      "source": [
        "# API service endpoint\n",
        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
        "\n",
        "# Vertex location root path for your dataset, model and endpoint resources\n",
        "PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + REGION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,prediction,cpu"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "Set the hardware accelerators (e.g., GPU), if any, for training and prediction.\n",
        "\n",
        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aip.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "For GPU, available accelerators include:\n",
        "   - aip.AcceleratorType.NVIDIA_TESLA_K80\n",
        "   - aip.AcceleratorType.NVIDIA_TESLA_P100\n",
        "   - aip.AcceleratorType.NVIDIA_TESLA_P4\n",
        "   - aip.AcceleratorType.NVIDIA_TESLA_T4\n",
        "   - aip.AcceleratorType.NVIDIA_TESLA_V100\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3 -- which is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kVb-KlTfTA89"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
        "    TRAIN_GPU, TRAIN_NGPU = (\n",
        "        vertex_ai_v1beta1.AcceleratorType.NVIDIA_TESLA_K80,\n",
        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
        "    )\n",
        "else:\n",
        "    TRAIN_GPU, TRAIN_NGPU = (vertex_ai_v1beta1.AcceleratorType.NVIDIA_TESLA_K80, 1)\n",
        "\n",
        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
        "        vertex_ai_v1beta1.AcceleratorType.NVIDIA_TESLA_K80,\n",
        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
        "    )\n",
        "else:\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for training and prediction.\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers).\n",
        "\n",
        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CPty8w-cTBtZ"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TF\"):\n",
        "    TF = os.getenv(\"IS_TESTING_TF\")\n",
        "else:\n",
        "    TF = \"2-8\"\n",
        "\n",
        "if TF[0] == \"2\":\n",
        "    if TRAIN_GPU:\n",
        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "    if DEPLOY_GPU:\n",
        "        DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "else:\n",
        "    if TRAIN_GPU:\n",
        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "    if DEPLOY_GPU:\n",
        "        DEPLOY_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training,prediction"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training and prediction.\n",
        "\n",
        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you will use for for training and prediction.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: The following is not supported for training:*\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cg2J5hFJOzvP"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_MACHINE\"):\n",
        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_TRAIN_MACHINE\")\n",
        "else:\n",
        "    MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)\n",
        "\n",
        "if os.getenv(\"IS_TESTING_DEPLOY_MACHINE\"):\n",
        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_DEPLOY_MACHINE\")\n",
        "else:\n",
        "    MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "DEPLOY_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hMY97zSZLVj7"
      },
      "source": [
        "### Helpers\n",
        "\n",
        "Below a series of helper functions you use along the tutorial:\n",
        "\n",
        "\n",
        "\n",
        "1.   `create_index_to_name_map` : Creates a map from index to name for label variables\n",
        "2.   `extract_images_and_labels` : A function to extract a batch of images and labels from a dataset\n",
        "3.   `plot_input_and_neighbors` : A function to plot the input image and its neighbors\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tuhbXWKrLYuk"
      },
      "outputs": [],
      "source": [
        "def create_index_to_name_map(ds_info):\n",
        "    \"\"\"\n",
        "    Creates a map from label name to numerical index.\n",
        "    Args:\n",
        "        ds_info: DatasetInfo object.\n",
        "    Returns:\n",
        "        index_to_name_map: dict. Map from name to index.\n",
        "    \"\"\"\n",
        "    index_to_name = {}\n",
        "    num_classes = ds_info.features[\"label\"].num_classes\n",
        "    names = ds_info.features[\"label\"].names\n",
        "    for i in range(num_classes):\n",
        "        index_to_name[i] = names[i]\n",
        "    return index_to_name\n",
        "\n",
        "\n",
        "def extract_images_and_labels(ds, num_batches):\n",
        "    \"\"\"\n",
        "    Extract images and labels from a dataset.\n",
        "    Args:\n",
        "        ds: A dataset.\n",
        "        num_batches: The number of batches to extract. -1 uses the whole dataset\n",
        "    Returns:\n",
        "        images: A numpy structure of images.\n",
        "        labels: A numpy structure of labels.\n",
        "    \"\"\"\n",
        "    data_slice = ds.take(num_batches)\n",
        "    images = []\n",
        "    labels = []\n",
        "    for image, label in data_slice:\n",
        "        images.append(image)\n",
        "        labels.append(label)\n",
        "    images = tf.concat(images, 0)\n",
        "    labels = tf.concat(labels, 0)\n",
        "    print(f\"Image batch shape: {images.shape}\")\n",
        "    return images.numpy(), labels.numpy()\n",
        "\n",
        "\n",
        "def plot_input_and_neighbors(\n",
        "    val_img_idx,\n",
        "    all_train_images,\n",
        "    val_images,\n",
        "    all_train_labels,\n",
        "    val_labels,\n",
        "    label_index_to_name,\n",
        "    data_with_neighbors,\n",
        "):\n",
        "    \"\"\"\n",
        "    Plot the input image and its neighbors.\n",
        "    Args:\n",
        "        val_img_idx: Index of the input image.\n",
        "        all_train_images: All training images.\n",
        "        val_images: Validation images.\n",
        "        all_train_labels: All training labels.\n",
        "        val_labels: Validation labels.\n",
        "        label_index_to_name: Dictionary mapping label indices to names.\n",
        "        data_with_neighbors: Data with neighbors.\n",
        "    Returns:\n",
        "        None\n",
        "    \"\"\"\n",
        "    image = val_images[val_img_idx]\n",
        "    fig = plt.figure(figsize=(24, 12))\n",
        "    ax_list = fig.subplots(3, 5)\n",
        "    ax_list[0, 0].axis(\"off\")\n",
        "    ax_list[0, 1].axis(\"off\")\n",
        "    ax_list[0, 3].axis(\"off\")\n",
        "    ax_list[0, 4].axis(\"off\")\n",
        "    ax = ax_list[0, 2]\n",
        "    class_label = val_labels[val_img_idx]\n",
        "    ax.set_title(\n",
        "        f\"{class_label}:{label_index_to_name[class_label]} (example index: {val_img_idx})\",\n",
        "        fontsize=15,\n",
        "    )\n",
        "    ax.axis(\"off\")\n",
        "    ax.imshow(image.astype(\"uint8\"))\n",
        "\n",
        "    neighbor_list = data_with_neighbors[val_img_idx][\"neighbors\"]\n",
        "    num_neighbors = len(neighbor_list)\n",
        "    for n in range(num_neighbors):\n",
        "        neighbor = neighbor_list[n]\n",
        "        neighbor_idx = int(neighbor[\"neighborId\"])\n",
        "        neighbor_dist = neighbor[\"neighborDistance\"]\n",
        "        ax = ax_list[1 + n // 5, n % 5]\n",
        "        class_label = all_train_labels[neighbor_idx]\n",
        "        ax.set_title(\n",
        "            f\"{class_label}:{label_index_to_name[class_label]} (dist: {neighbor_dist:.3f})\",\n",
        "            fontsize=15,\n",
        "        )\n",
        "        ax.axis(\"off\")\n",
        "        ax.imshow(all_train_images[neighbor_idx].astype(\"uint8\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gd6RoZHJYvt8"
      },
      "source": [
        "# Tutorial\n",
        "\n",
        "The Vertex Explainable Example-based API provides an highly performant ANN service for returning similar examples to new predictions/instances.\n",
        "\n",
        "Then, in order to leverage the Example-based explanations service, you need to cover the following steps:\n",
        "\n",
        "1) Index the entire dataset: You need to provide a path to an embedding model in a GCS bucket, training data stored in a GCS bucket and the config file for example-based explanation \n",
        "   \n",
        "2) Deploy index and model: You need to specify the machine to use and the model identifier from the model upload set\n",
        "\n",
        "3) Query for similar examples: You need to make the explain query and model will return similar examples\n",
        "\n",
        "Let's start to experiment a custom model for STL10 dataset you will use to extract embeddings. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-pUOiVFsZbm4"
      },
      "source": [
        "## Experimentation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QDKXS5h2LpNB"
      },
      "source": [
        "### Prepare training data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cDDCBHJbgW7n"
      },
      "source": [
        "#### Download and visualize the data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KGiCMhPeSYdy"
      },
      "outputs": [],
      "source": [
        "split_ds, ds_info = tfds.load(\n",
        "    DATASET_NAME,\n",
        "    split=[\"train\", \"test\"],\n",
        "    as_supervised=True,  # Include labels\n",
        "    with_info=True,\n",
        "    shuffle_files=False,  # ensuring that the data doesn't get shuffled between runs\n",
        "    data_dir=RAW_DIR,\n",
        ")\n",
        "train_ds, validation_ds = split_ds\n",
        "tfds.show_examples(ds=train_ds, ds_info=ds_info);"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sGNlBKlWju5t"
      },
      "outputs": [],
      "source": [
        "print(f'Number of classes in the dataset: {ds_info.features[\"label\"].num_classes}')\n",
        "print(f'Label names: {ds_info.features[\"label\"].names}')\n",
        "print(f'Number of examples in training split: {ds_info.splits[\"train\"].num_examples}')\n",
        "print(f'Number of examples in validation split: {ds_info.splits[\"test\"].num_examples}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IRmJEzcdhofK"
      },
      "source": [
        "#### Prepare the images\n",
        "\n",
        "Scale the images to the size expected by the downstream model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j3e_Cg4ljyuk"
      },
      "outputs": [],
      "source": [
        "size = (224, 224)\n",
        "train_ds = train_ds.map(lambda x, y: (tf.image.resize(x, size), y))\n",
        "validation_ds = validation_ds.map(lambda x, y: (tf.image.resize(x, size), y))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sjds2n3Bj-9w"
      },
      "source": [
        "#### Batch and prefetch data\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t1GjhJODkJRh"
      },
      "outputs": [],
      "source": [
        "batch_size = 32\n",
        "\n",
        "train_ds = train_ds.batch(batch_size).prefetch(buffer_size=10)\n",
        "validation_ds = validation_ds.batch(batch_size).prefetch(buffer_size=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vlhPZStokO90"
      },
      "outputs": [],
      "source": [
        "label_index_to_name = create_index_to_name_map(ds_info)\n",
        "print(label_index_to_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6xymTJVkdXa"
      },
      "source": [
        "### Create and train the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WPUESdjzhDWD"
      },
      "source": [
        "#### Fine-tune the last layer of a pre-trained classification model\n",
        "\n",
        "You use a `MobileNetV2` [Keras Application](https://keras.io/api/applications/) deep learning model that is available alongside pre-trained weights for fine-tuning the model which will be used to create embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lxnUTl_2i4oK"
      },
      "outputs": [],
      "source": [
        "# Each image can be flipped and rotated to generate more training data, and\n",
        "# ensure the model is more robust to such changes, since a rotated \"bird\" should\n",
        "# still be classified as a bird.\n",
        "\n",
        "data_augmentation = keras.Sequential(\n",
        "    [\n",
        "        layers.experimental.preprocessing.RandomFlip(\"horizontal\"),\n",
        "        layers.experimental.preprocessing.RandomRotation(0.1),\n",
        "    ]\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fUS5WlTei5BD"
      },
      "outputs": [],
      "source": [
        "size_3d = (224, 224, 3)\n",
        "num_classes = ds_info.features[\"label\"].num_classes\n",
        "\n",
        "# Load an ImageNet model.\n",
        "base_model = keras.applications.MobileNetV2(\n",
        "    weights=\"imagenet\",  # Load weights pre-trained on ImageNet.\n",
        "    input_shape=size_3d,\n",
        "    include_top=False,\n",
        ")  # Do not include the ImageNet classifier at the top.\n",
        "\n",
        "# Freeze the base_model\n",
        "base_model.trainable = False\n",
        "\n",
        "# Create new model on top\n",
        "inputs = keras.Input(shape=size_3d)\n",
        "x = data_augmentation(inputs)  # Apply random data augmentation\n",
        "\n",
        "# Pre-trained Xception weights requires that input be normalized\n",
        "# from (0, 255) to a range (-1., +1.), the normalization layer\n",
        "# does the following, outputs = (inputs - mean) / sqrt(var)\n",
        "norm_layer = keras.layers.experimental.preprocessing.Normalization()\n",
        "mean = np.array([127.5] * 3)\n",
        "var = mean**2\n",
        "\n",
        "# Scale inputs to [-1, +1]\n",
        "x = norm_layer(x)\n",
        "norm_layer.set_weights([mean, var, 0])\n",
        "\n",
        "# The base model contains batchnorm layers. We want to keep them in inference mode\n",
        "# when we unfreeze the base model for fine-tuning, so we make sure that the\n",
        "# base_model is running in inference mode here.\n",
        "x = base_model(x, training=False)\n",
        "x = keras.layers.GlobalAveragePooling2D()(x)\n",
        "x = keras.layers.Dropout(0.2)(x)  # Regularize with dropout\n",
        "outputs = keras.layers.Dense(num_classes)(x)\n",
        "model = keras.Model(inputs, outputs)\n",
        "\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KhhGgjqyeozX"
      },
      "source": [
        "You can improve the latent representation of the pre-trained model by unfreezing the lower layers and fine-tuning. However, you forego that step in this demo so that the learned representations are exactly the same as the pre-trained model, and only the softmax layer is trained briefly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rO0blIoOjihC"
      },
      "outputs": [],
      "source": [
        "model.compile(\n",
        "    optimizer=keras.optimizers.Adam(),\n",
        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
        ")\n",
        "\n",
        "epochs = 2\n",
        "history = model.fit(train_ds, epochs=epochs, validation_data=validation_ds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YkNdPtuNhofN"
      },
      "source": [
        "### Evaluate the model\n",
        "\n",
        "Plot the `accurancy` and the `loss` of the model after training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FxzsSYunlFGh"
      },
      "outputs": [],
      "source": [
        "x_axis = range(1, epochs + 1)\n",
        "plt.figure(figsize=(20, 8))\n",
        "plt.subplot(1, 2, 1)\n",
        "plt.plot(x_axis, history.history[\"sparse_categorical_accuracy\"])\n",
        "plt.plot(x_axis, history.history[\"val_sparse_categorical_accuracy\"])\n",
        "plt.title(\"Model Accuracy\", fontsize=20)\n",
        "plt.ylabel(\"Accuracy\", fontsize=15)\n",
        "plt.xlabel(\"Epoch\", fontsize=15)\n",
        "plt.legend([\"Train\", \"Test\"], fontsize=15)\n",
        "plt.grid()\n",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.plot(x_axis, history.history[\"loss\"])\n",
        "plt.plot(x_axis, history.history[\"val_loss\"])\n",
        "plt.title(\"Model Loss\", fontsize=20)\n",
        "plt.ylabel(\"Loss\", fontsize=15)\n",
        "plt.xlabel(\"Epoch\", fontsize=15)\n",
        "plt.legend([\"Train\", \"Test\"], fontsize=15)\n",
        "plt.grid();"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "59OE0kYjaqe5"
      },
      "source": [
        "## Formalization\n",
        "\n",
        "Now that you experiment the model to create the embeddings, let's formalize the training to leverage Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MAXCt-Ldplv2"
      },
      "source": [
        "### Set up clients\n",
        "\n",
        "The Vertex AI client library works as a client/server model. Then you need to set clients to use different services.\n",
        "\n",
        "You will use different clients in this tutorial for different steps in the workflow. So set them all up upfront.\n",
        "\n",
        "- Model Service for `Model` resources.\n",
        "- Endpoint Service for deployment.\n",
        "- Job Service for batch jobs and custom training.\n",
        "- Prediction Service for serving."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vq5UsPfqPiKi"
      },
      "outputs": [],
      "source": [
        "# client options same for all services\n",
        "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
        "\n",
        "\n",
        "def create_job_client():\n",
        "    client = vertex_ai_v1beta1.JobServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "def create_model_client():\n",
        "    client = vertex_ai_v1beta1.ModelServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "def create_endpoint_client():\n",
        "    client = vertex_ai_v1beta1.EndpointServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "def create_prediction_client():\n",
        "    client = vertex_ai_v1beta1.PredictionServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "clients = {}\n",
        "clients[\"job\"] = create_job_client()\n",
        "clients[\"model\"] = create_model_client()\n",
        "clients[\"endpoint\"] = create_endpoint_client()\n",
        "clients[\"prediction\"] = create_prediction_client()\n",
        "\n",
        "for client in clients.items():\n",
        "    print(client)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_model"
      },
      "source": [
        "### Train a model\n",
        "\n",
        "There are two ways you can train a custom model using a container image:\n",
        "\n",
        "- **Use a Google Cloud prebuilt container**. If you use a prebuilt container, you will additionally specify a Python package to install into the container image. This Python package contains your code for training a custom model.\n",
        "\n",
        "- **Use your own custom container image**. If you use your own container, the container needs to contain your code for training a custom model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_specification:prebuilt_container"
      },
      "source": [
        "### Prepare your custom job specification\n",
        "\n",
        "Now that your clients are ready, your first step is to create a Job Specification for your custom training job. The job specification will consist of the following:\n",
        "\n",
        "- `worker_pool_spec` : The specification of the type of machine(s) you will use for training and how many (single or distributed)\n",
        "- `python_package_spec` : The specification of the Python package to be installed with the pre-built container."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_machine_specification"
      },
      "source": [
        "#### Prepare your machine specification\n",
        "\n",
        "Now define the machine specification for your custom training job. This tells Vertex AI what type of machine instance to provision for the training.\n",
        "  - `machine_type`: The type of GCP instance to provision -- e.g., n1-standard-8.\n",
        "  - `accelerator_type`: The type, if any, of hardware accelerator. In this tutorial if you previously set the variable `TRAIN_GPU != None`, you are using a GPU; otherwise you will use a CPU.\n",
        "  - `accelerator_count`: The number of accelerators."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u5hIFmP4OzvR"
      },
      "outputs": [],
      "source": [
        "if TRAIN_GPU:\n",
        "    machine_spec = {\n",
        "        \"machine_type\": TRAIN_COMPUTE,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "    }\n",
        "else:\n",
        "    machine_spec = {\"machine_type\": TRAIN_COMPUTE, \"accelerator_count\": 0}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_disk_specification"
      },
      "source": [
        "#### Prepare your disk specification\n",
        "\n",
        "(optional) Now define the disk specification for your custom training job. This tells Vertex what type and size of disk to provision in each machine instance for the training.\n",
        "\n",
        "  - `boot_disk_type`: Either SSD or Standard. SSD is faster, and Standard is less expensive. Defaults to SSD.\n",
        "  - `boot_disk_size_gb`: Size of disk in GB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1MfgWif9OzvS"
      },
      "outputs": [],
      "source": [
        "DISK_TYPE = \"pd-ssd\"  # [ pd-ssd, pd-standard]\n",
        "DISK_SIZE = 200  # GB\n",
        "\n",
        "disk_spec = {\"boot_disk_type\": DISK_TYPE, \"boot_disk_size_gb\": DISK_SIZE}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_worker_pool_specification:prebuilt_container"
      },
      "source": [
        "#### Define the worker pool specification\n",
        "\n",
        "Next, you define the worker pool specification for your custom training job. The worker pool specification will consist of the following:\n",
        "\n",
        "- `replica_count`: The number of instances to provision of this machine type.\n",
        "- `machine_spec`: The hardware specification.\n",
        "- `disk_spec` : (optional) The disk storage specification.\n",
        "\n",
        "- `python_package`: The Python training package to install on the VM instance(s) and which Python module to invoke, along with command line arguments for the Python module.\n",
        "\n",
        "Let's dive deeper now into the python package specification:\n",
        "\n",
        "-`executor_image_spec`: This is the docker image which is configured for your custom training job.\n",
        "\n",
        "-`package_uris`: This is a list of the locations (URIs) of your python training packages to install on the provisioned instance. The locations need to be in a Cloud Storage bucket. These can be either individual python files or a zip (archive) of an entire package. In the later case, the job service will unzip (unarchive) the contents into the docker image.\n",
        "\n",
        "-`python_module`: The Python module (script) to invoke for running the custom training job. In this example, you will be invoking `trainer.task.py` -- note that it was not neccessary to append the `.py` suffix.\n",
        "\n",
        "-`args`: The command line arguments to pass to the corresponding Pythom module. In this example, you will be setting:\n",
        "  - `\"--model-dir=\" + MODEL_URI` : The Cloud Storage location where to store the model artifacts. There are two ways to tell the training script where to save the model artifacts:\n",
        "      - direct: You pass the Cloud Storage location as a command line argument to your training script (set variable `DIRECT = True`), or\n",
        "      - indirect: The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script (set variable `DIRECT = False`). In this case, you tell the service the model artifact location in the job specification.\n",
        "  - `\"--epochs=\" + EPOCHS`: The number of epochs for training.\n",
        "  - `\"--distribute=\" + TRAIN_STRATEGY\"` : The training distribution strategy to use for single or distributed training.\n",
        "     - `\"single\"`: single device.\n",
        "     - `\"mirror\"`: all GPU devices on a single compute instance.\n",
        "     - `\"multi\"`: all GPU devices on all compute instances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJCovykJOzvS"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = \"custom_job_\" + TIMESTAMP\n",
        "MODEL_URI = f\"{BUCKET_URI}/{DELIVERABLE_PATH}/models/mobilenetv2-{DATASET_NAME}\"\n",
        "\n",
        "if not TRAIN_NGPU or TRAIN_NGPU < 2:\n",
        "    TRAIN_STRATEGY = \"single\"\n",
        "else:\n",
        "    TRAIN_STRATEGY = \"mirror\"\n",
        "\n",
        "EPOCHS = 2\n",
        "\n",
        "DIRECT = True\n",
        "if DIRECT:\n",
        "    CMDARGS = [\n",
        "        \"--model-dir=\" + MODEL_URI,\n",
        "        \"--epochs=\" + str(EPOCHS),\n",
        "        \"--distribute=\" + TRAIN_STRATEGY,\n",
        "    ]\n",
        "else:\n",
        "    CMDARGS = [\n",
        "        \"--epochs=\" + str(EPOCHS),\n",
        "        \"--distribute=\" + TRAIN_STRATEGY,\n",
        "    ]\n",
        "\n",
        "worker_pool_spec = [\n",
        "    {\n",
        "        \"replica_count\": 1,\n",
        "        \"machine_spec\": machine_spec,\n",
        "        \"disk_spec\": disk_spec,\n",
        "        \"python_package_spec\": {\n",
        "            \"executor_image_uri\": TRAIN_IMAGE,\n",
        "            \"package_uris\": [BUCKET_URI + \"/trainer_stl10.tar.gz\"],\n",
        "            \"python_module\": \"trainer.task\",\n",
        "            \"args\": CMDARGS,\n",
        "        },\n",
        "    }\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "assemble_custom_job_specification"
      },
      "source": [
        "### Assemble a job specification\n",
        "\n",
        "Now assemble the complete description for the custom job specification:\n",
        "\n",
        "- `display_name`: The human readable name you assign to this custom job.\n",
        "- `job_spec`: The specification for the custom job.\n",
        "    - `worker_pool_specs`: The specification for the machine VM instances.\n",
        "    - `base_output_directory`: This tells the service the Cloud Storage location where to save the model artifacts (when variable `DIRECT = False`). The service will then pass the location to the training script as the environment variable `AIP_MODEL_DIR`, and the path will be of the form:\n",
        "\n",
        "                <output_uri_prefix>/model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qoY_r-2WOzvS"
      },
      "outputs": [],
      "source": [
        "if DIRECT:\n",
        "    job_spec = {\"worker_pool_specs\": worker_pool_spec}\n",
        "else:\n",
        "    job_spec = {\n",
        "        \"worker_pool_specs\": worker_pool_spec,\n",
        "        \"base_output_directory\": {\"output_uri_prefix\": MODEL_DIR},\n",
        "    }\n",
        "\n",
        "custom_job = {\"display_name\": JOB_NAME, \"job_spec\": job_spec}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package"
      },
      "source": [
        "### Examine the training package\n",
        "\n",
        "#### Package layout\n",
        "\n",
        "Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
        "In the following cells, you will assemble the training package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iNcSkITGOzvT"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "! rm -rf custom\n",
        "! mkdir custom\n",
        "\n",
        "# Add package information\n",
        "! touch custom/README.md\n",
        "\n",
        "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
        "! echo \"$setup_cfg\" > custom/setup.cfg\n",
        "\n",
        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'tensorflow_datasets==4.0.1',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
        "! echo \"$setup_py\" > custom/setup.py\n",
        "\n",
        "pkg_info = \"Metadata-Version: 1.0\\n\\nName: STL10 image classification\\n\\nVersion: 0.0.0\\n\\nSummary: Demonstration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: googler@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
        "! echo \"$pkg_info\" > custom/PKG-INFO\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir custom/trainer\n",
        "! touch custom/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taskpy_contents:cifar10"
      },
      "source": [
        "#### Create the task script\n",
        "\n",
        "In the next cell, you write the contents of the training script task.py. We won't go into detail, it's just there for you to browse. In summary:\n",
        "\n",
        "- Get the directory where to save the model artifacts from the command line (`--model_dir`), and if not specified, then from the environment variable `AIP_MODEL_DIR`.\n",
        "- Loads STL10 dataset from TF Datasets (tfds).\n",
        "- Builds a model using TF.Keras model API.\n",
        "- Compiles the model (`compile()`).\n",
        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
        "- Trains the model (`fit()`) with epochs and steps according to the arguments `args.epochs`\n",
        "- Saves the trained model (`save(args.model_dir)`) to the specified model directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YogsgSLpOzvT"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/trainer/task.py\n",
        "# Single, Mirror and Multi-Machine Distributed Training for STL10\n",
        "\n",
        "# General\n",
        "import os\n",
        "import argparse\n",
        "import sys\n",
        "import logging\n",
        "\n",
        "logging.basicConfig(format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')\n",
        "\n",
        "# Training\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "from tensorflow.python.client import device_lib\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "\n",
        "# Variables\n",
        "N_CLASSES = 10\n",
        "SIZE = (224, 224)\n",
        "BUFFER_SIZE = 10\n",
        "BATCH_SIZE = 32\n",
        "SIZE_3D = (224, 224, 3)\n",
        "\n",
        "\n",
        "# Helpers\n",
        "# Get arguments\n",
        "def get_args():\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument('--model-dir', dest='model_dir',\n",
        "                        default=os.getenv(\"AIP_MODEL_DIR\"), type=str, help='Model dir.')\n",
        "    parser.add_argument('--epochs', dest='epochs',\n",
        "                        default=2, type=int,\n",
        "                        help='Number of epochs.')\n",
        "    parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
        "                        help='distributed training strategy')\n",
        "    args = parser.parse_args()\n",
        "    return args\n",
        "\n",
        "\n",
        "# Single Machine, single compute device\n",
        "def get_strategy(distribute):\n",
        "    if distribute == 'single':\n",
        "        if tf.config.list_physical_devices('GPU'):\n",
        "            strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
        "        else:\n",
        "            strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "\n",
        "    # Single Machine, multiple compute device\n",
        "    elif distribute == 'mirror':\n",
        "        strategy = tf.distribute.MirroredStrategy()\n",
        "\n",
        "    # Multiple Machine, multiple compute device\n",
        "    elif distribute == 'multi':\n",
        "        strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\n",
        "    return strategy\n",
        "\n",
        "\n",
        "# Get train and validation datasets\n",
        "def get_train_val_ds():\n",
        "    split_ds = tfds.load(\n",
        "        'stl10',\n",
        "        split=[\"train\", \"test\"],\n",
        "        as_supervised=True,\n",
        "        shuffle_files=False,\n",
        "    )\n",
        "    train_ds, validation_ds = split_ds\n",
        "    return train_ds, validation_ds\n",
        "\n",
        "\n",
        "# Get preprocessed dataset\n",
        "def preprocess_dataset(train_ds, validation_ds):\n",
        "    train_ds = train_ds.map(lambda x, y: (tf.image.resize(x, SIZE), y))\n",
        "    validation_ds = validation_ds.map(lambda x, y: (tf.image.resize(x, SIZE), y))\n",
        "    return train_ds, validation_ds\n",
        "\n",
        "\n",
        "# Build the model\n",
        "def build_and_compile_model():\n",
        "    # Define data augmentation layers\n",
        "    data_augmentation = keras.Sequential(\n",
        "        [\n",
        "            layers.experimental.preprocessing.RandomFlip(\"horizontal\"),\n",
        "            layers.experimental.preprocessing.RandomRotation(0.1),\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    # Define base model\n",
        "    base_model = keras.applications.MobileNetV2(\n",
        "        weights=\"imagenet\",\n",
        "        input_shape=SIZE_3D,\n",
        "        include_top=False,\n",
        "    )\n",
        "\n",
        "    # Freeze the base_model\n",
        "    base_model.trainable = False\n",
        "\n",
        "    # Create new model on top\n",
        "    inputs = keras.Input(shape=SIZE_3D)\n",
        "    x = data_augmentation(inputs)\n",
        "\n",
        "    # Normalize inputs for Pre-trained Xception weights\n",
        "    norm_layer = keras.layers.experimental.preprocessing.Normalization()\n",
        "    mean = np.array([127.5] * 3)\n",
        "    var = mean ** 2\n",
        "    x = norm_layer(x)\n",
        "    norm_layer.set_weights([mean, var, 0])\n",
        "\n",
        "    # Set inference mode\n",
        "    x = base_model(x, training=False)\n",
        "    x = keras.layers.GlobalAveragePooling2D()(x)\n",
        "    x = keras.layers.Dropout(0.2)(x)  # Regularize with dropout\n",
        "    outputs = keras.layers.Dense(N_CLASSES)(x)\n",
        "\n",
        "    # Build the model\n",
        "    model = keras.Model(inputs, outputs)\n",
        "    model.summary()\n",
        "\n",
        "    # Compile the model \n",
        "    model.compile(\n",
        "        optimizer=keras.optimizers.Adam(),\n",
        "        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "        metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
        "    )\n",
        "\n",
        "    return model\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "\n",
        "    # Initialize arguments and strategy\n",
        "    logging.info('Initialize arguments and strategy')\n",
        "    train_args = get_args()\n",
        "\n",
        "    # Variables\n",
        "    strategy = get_strategy(train_args.distribute)\n",
        "    epochs = train_args.epochs\n",
        "    model_dir = train_args.model_dir\n",
        "    n_workers = strategy.num_replicas_in_sync\n",
        "    global_batch_size = BATCH_SIZE * n_workers\n",
        "\n",
        "    # Print versions\n",
        "    logging.info('Python Version = {}'.format(sys.version))\n",
        "    logging.info('TensorFlow Version = {}'.format(tf.__version__))\n",
        "    logging.info('Python Version = {}'.format(sys.version))\n",
        "\n",
        "    # Get TRAIN and VAL datasets\n",
        "    logging.info('Get train and val datasets')\n",
        "    train_ds, validation_ds = get_train_val_ds()\n",
        "\n",
        "    # Preprocess TRAIN and VAL datasets\n",
        "    logging.info('Get train and val datasets')\n",
        "    train_ds, validation_ds = preprocess_dataset(train_ds, validation_ds)\n",
        "    train_ds = train_ds.batch(global_batch_size).prefetch(buffer_size=BUFFER_SIZE)\n",
        "    validation_ds = validation_ds.batch(global_batch_size).prefetch(buffer_size=BUFFER_SIZE)\n",
        "\n",
        "    # Train model\n",
        "    logging.info('Build model')\n",
        "    with strategy.scope():\n",
        "        model = build_and_compile_model()\n",
        "\n",
        "    logging.info('Train model')\n",
        "    model.fit(train_ds, epochs=epochs, validation_data=validation_ds)\n",
        "    logging.info(f'Save model in {model_dir}')\n",
        "    model.save(model_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tarball_training_script"
      },
      "source": [
        "#### Store training script on your Cloud Storage bucket\n",
        "\n",
        "Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zEx3hg05OzvT"
      },
      "outputs": [],
      "source": [
        "! rm -f custom.tar custom.tar.gz\n",
        "! tar cvf custom.tar custom\n",
        "! gzip custom.tar\n",
        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_stl10.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job"
      },
      "source": [
        "### Train the model \n",
        "\n",
        "\n",
        "Now start the training of your custom training job on Vertex. Use this helper function `create_custom_job`, which takes the following parameter:\n",
        "\n",
        "-`custom_job`: The specification for the custom job.\n",
        "\n",
        "The helper function calls job client service's `create_custom_job` method, with the following parameters:\n",
        "\n",
        "-`parent`: The Vertex location path to `Dataset`, `Model` and `Endpoint` resources.\n",
        "-`custom_job`: The specification for the custom job.\n",
        "\n",
        "You will display a handful of the fields returned in `response` object, with the two that are of most interest are:\n",
        "\n",
        "`response.name`: The Vertex fully qualified identifier assigned to this custom training job. You save this identifier for using in subsequent steps.\n",
        "\n",
        "`response.state`: The current state of the custom training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZfcnibeAOzvU"
      },
      "outputs": [],
      "source": [
        "def create_custom_job(custom_job):\n",
        "    response = clients[\"job\"].create_custom_job(parent=PARENT, custom_job=custom_job)\n",
        "    print(\"name:\", response.name)\n",
        "    print(\"display_name:\", response.display_name)\n",
        "    print(\"state:\", response.state)\n",
        "    print(\"create_time:\", response.create_time)\n",
        "    print(\"update_time:\", response.update_time)\n",
        "    return response\n",
        "\n",
        "\n",
        "response = create_custom_job(custom_job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "job_id:response"
      },
      "source": [
        "Now get the unique identifier for the custom job you created."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "goXUwSIQOzvU"
      },
      "outputs": [],
      "source": [
        "# The full unique ID for the custom job\n",
        "job_id = response.name\n",
        "# The short numeric ID for the custom job\n",
        "job_short_id = job_id.split(\"/\")[-1]\n",
        "\n",
        "print(job_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "get_custom_job"
      },
      "source": [
        "### Get information on a custom job\n",
        "\n",
        "Next, use this helper function `get_custom_job`, which takes the following parameter:\n",
        "\n",
        "- `name`: The Vertex fully qualified identifier for the custom job.\n",
        "\n",
        "The helper function calls the job client service's`get_custom_job` method, with the following parameter:\n",
        "\n",
        "- `name`: The Vertex fully qualified identifier for the custom job.\n",
        "\n",
        "If you recall, you got the Vertex fully qualified identifier for the custom job in the `response.name` field when you called the `create_custom_job` method, and saved the identifier in the variable `job_id`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NecLamxVOzvV"
      },
      "outputs": [],
      "source": [
        "def get_custom_job(name, silent=False):\n",
        "    response = clients[\"job\"].get_custom_job(name=name)\n",
        "    if silent:\n",
        "        return response\n",
        "\n",
        "    print(\"name:\", response.name)\n",
        "    print(\"display_name:\", response.display_name)\n",
        "    print(\"state:\", response.state)\n",
        "    print(\"create_time:\", response.create_time)\n",
        "    print(\"update_time:\", response.update_time)\n",
        "    return response\n",
        "\n",
        "\n",
        "response = get_custom_job(job_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ijuYBvLOzvW"
      },
      "outputs": [],
      "source": [
        "while True:\n",
        "    response = get_custom_job(job_id, True)\n",
        "    if response.state != vertex_ai_v1beta1.JobState.JOB_STATE_SUCCEEDED:\n",
        "        print(\"Training job has not completed:\", response.state)\n",
        "        model_path_to_deploy = None\n",
        "        if response.state == vertex_ai_v1beta1.JobState.JOB_STATE_FAILED:\n",
        "            break\n",
        "    else:\n",
        "        if not DIRECT:\n",
        "            MODEL_DIR = MODEL_DIR + \"/model\"\n",
        "        model_path_to_deploy = MODEL_URI\n",
        "        print(\"Training Time:\", response.update_time - response.create_time)\n",
        "        break\n",
        "    time.sleep(60)\n",
        "\n",
        "print(\"model_to_deploy:\", model_path_to_deploy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C6jqOH23PjGK"
      },
      "source": [
        "## Deployment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "load_saved_model"
      },
      "source": [
        "### Load the saved model\n",
        "\n",
        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now load it from the Cloud Storage bucket, and then you can do some things, like evaluate the model, and do a prediction.\n",
        "\n",
        "To load, you use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6a8NKCrJOzvW"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.models.load_model(MODEL_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "how_serving_function_works"
      },
      "source": [
        "### Prepare the model for serving\n",
        "\n",
        "Next, you will upload your TF.Keras model from the custom job to Vertex AI `Model` service, which will create a Vertex `Model` resource for your custom model. During upload, you need to define a serving function to convert data to the format your model expects. If you send encoded data to Vertex AI, your serving function ensures that the data is decoded on the model server before it is passed as input to your model.\n",
        "\n",
        "### How does the serving function work\n",
        "\n",
        "When you send a request to an online prediction server, the request is received by a HTTP server. The HTTP server extracts the prediction request from the HTTP request content body. The extracted prediction request is forwarded to the serving function. For Google pre-built prediction containers, the request content is passed to the serving function as a `tf.string`.\n",
        "\n",
        "The serving function consists of two parts:\n",
        "\n",
        "- `preprocessing function`:\n",
        "  - Converts the input (`tf.string`) to the input shape and data type of the underlying model (dynamic graph).\n",
        "  - Performs the same preprocessing of the data that was done during training the underlying model -- e.g., normalizing, scaling, etc.\n",
        "- `post-processing function`:\n",
        "  - Converts the model output to format expected by the receiving application -- e.q., compresses the output.\n",
        "  - Packages the output for the the receiving application -- e.g., add headings, make JSON object, etc.\n",
        "\n",
        "Both the preprocessing and post-processing functions are converted to static graphs which are fused to the model. The output from the underlying model is passed to the post-processing function. The post-processing function passes the converted/packaged output back to the HTTP server. The HTTP server returns the output as the HTTP response content.\n",
        "\n",
        "One consideration you need to consider when building serving functions for TF.Keras models is that they run as static graphs. That means, you cannot use TF graph operations that require a dynamic graph. If you do, you will get an error during the compile of the serving function which will indicate that you are using an EagerTensor which is not supported."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "serving_function_image:xai"
      },
      "source": [
        "#### Create serving function for image data\n",
        "\n",
        "To pass images to the prediction service, you encode the compressed (e.g., JPEG) image bytes into base 64 -- which makes the content safe from modification while transmitting binary data over the network. Since this deployed model expects input data as raw (uncompressed) bytes, you need to ensure that the base 64 encoded data gets converted back to raw bytes before it is passed as input to the deployed model.\n",
        "\n",
        "To resolve this, define a serving function (`serving_fn`) and attach it to the model as a preprocessing step. Add a `@tf.function` decorator so the serving function is fused to the underlying model (instead of upstream on a CPU).\n",
        "\n",
        "When you send a prediction or explanation request, the content of the request is base 64 decoded into a Tensorflow string (`tf.string`), which is passed to the serving function (`serving_fn`). The serving function preprocesses the `tf.string` into raw (uncompressed) numpy bytes (`preprocess_fn`) to match the input requirements of the model:\n",
        "- `io.decode_jpeg`- Decompresses the JPG image which is returned as a Tensorflow tensor with three channels (RGB).\n",
        "- `image.convert_image_dtype` - Changes integer pixel values to float 32, and rescales pixel data between 0 and 1.\n",
        "- `image.resize` - Resizes the image to match the input shape for the model.\n",
        "\n",
        "At this point, the data can be passed to the model (`m_call`).\n",
        "\n",
        "#### XAI signatures\n",
        "\n",
        "When the serving function is saved back with the underlying model (`tf.saved_model.save`), you specify the input layer of the serving function as the signature `serving_default`.\n",
        "\n",
        "For XAI image models, you need to save two additional signatures from the serving function:\n",
        "\n",
        "- `xai_preprocess`: The preprocessing function in the serving function.\n",
        "- `xai_model`: The concrete function for calling the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nwoYEzGLOzvX"
      },
      "outputs": [],
      "source": [
        "CONCRETE_INPUT = \"numpy_inputs\"\n",
        "\n",
        "\n",
        "def _preprocess(bytes_input):\n",
        "    \"\"\"\n",
        "    The preprocess function.\n",
        "    Args:\n",
        "        bytes_input: The input image in bytes.\n",
        "    Returns:\n",
        "        The preprocessed image in numpy array.\n",
        "    \"\"\"\n",
        "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
        "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
        "    resized = tf.image.resize(decoded, size=(224, 224))\n",
        "    rescale = tf.cast(resized, tf.float32)\n",
        "    return rescale\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def preprocess_fn(bytes_inputs):\n",
        "    \"\"\"\n",
        "    Preprocess the input image.\n",
        "    Args:\n",
        "        bytes_inputs: A list of raw image bytes.\n",
        "    Returns:\n",
        "        A list of preprocessed images.\n",
        "    \"\"\"\n",
        "    decoded_images = tf.nest.map_structure(\n",
        "        tf.stop_gradient, tf.map_fn(_preprocess, bytes_inputs, dtype=tf.float32)\n",
        "    )\n",
        "    return {CONCRETE_INPUT: decoded_images}\n",
        "\n",
        "\n",
        "@tf.function(\n",
        "    input_signature=[tf.TensorSpec([None], tf.string), tf.TensorSpec([None], tf.string)]\n",
        ")\n",
        "def serving_fn(id, bytes_inputs):\n",
        "    \"\"\"\n",
        "    This function is used to serve the embeddings.\n",
        "    Args:\n",
        "        id: The id of the input.\n",
        "        bytes_inputs: The input image.\n",
        "    Returns:\n",
        "        The output of the model.\n",
        "    \"\"\"\n",
        "    images = preprocess_fn(bytes_inputs)\n",
        "    embedding = m_call(**images)\n",
        "    return {\"id\": id, \"embedding\": embedding}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1K414cl3hlME"
      },
      "source": [
        "#### Extract and upload the embedding model for index creation\n",
        "\n",
        "As we mentioned above, you need to provide embeddings in order to index the dataset. In this case, you skip the data augmentation layer and drop the softmax layer to get to the embeddings from the model you previously trained. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qNeutW2HmQW5"
      },
      "outputs": [],
      "source": [
        "embedding_model = keras.Sequential()\n",
        "for layer in model.layers[:-1]:  # go through until last layer\n",
        "    print(layer.name)\n",
        "    if \"sequential\" not in layer.name:  # skip data augmentation layer\n",
        "        embedding_model.add(layer)\n",
        "embedding_model.summary()\n",
        "probability_model = keras.Sequential([model, tf.keras.layers.Softmax()])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O7ZQqOlbnpKp"
      },
      "outputs": [],
      "source": [
        "EMBEDDINGS_URI = (\n",
        "    f\"{BUCKET_URI}/{DELIVERABLE_PATH}/embeddings/mobilenetv2-{DATASET_NAME}\"\n",
        ")\n",
        "\n",
        "m_call = tf.function(embedding_model.call).get_concrete_function(\n",
        "    [tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32, name=CONCRETE_INPUT)]\n",
        ")\n",
        "\n",
        "tf.saved_model.save(\n",
        "    embedding_model,\n",
        "    EMBEDDINGS_URI,\n",
        "    signatures={\n",
        "        \"serving_default\": serving_fn,\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "serving_function_signature:xai"
      },
      "source": [
        "### Get the serving function signature\n",
        "\n",
        "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
        "\n",
        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you will use later when you make a prediction request.\n",
        "\n",
        "You also need to know the name of the serving function's input and output layer for constructing the explanation metadata -- which is discussed subsequently."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9rLm9QUkOzvY"
      },
      "outputs": [],
      "source": [
        "embedding_model_loaded = tf.saved_model.load(EMBEDDINGS_URI)\n",
        "\n",
        "serving_input = list(\n",
        "    embedding_model_loaded.signatures[\"serving_default\"]\n",
        "    .structured_input_signature[1]\n",
        "    .keys()\n",
        ")[0]\n",
        "print(\"Serving function input:\", serving_input)\n",
        "serving_output = list(\n",
        "    embedding_model_loaded.signatures[\"serving_default\"].structured_outputs.keys()\n",
        ")[0]\n",
        "print(\"Serving function output:\", serving_output)\n",
        "\n",
        "input_name = model.input.name\n",
        "print(\"Model input name:\", input_name)\n",
        "output_name = model.output.name\n",
        "print(\"Model output name:\", output_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f4Y9WltYOwIE"
      },
      "source": [
        "### Upload the training data\n",
        "\n",
        "Next step is uploading the training data. The explanations are picked from this data. You can choose a smaller number of batches for a faster run but the results would be less precise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KrKfZB_xw0Qu"
      },
      "outputs": [],
      "source": [
        "if not os.path.exists(PREPROCESSED_DIR):\n",
        "    os.makedirs(PREPROCESSED_DIR, exist_ok=True)\n",
        "\n",
        "dataset_file = f\"{DATASET_NAME}-train-images.jsonl\"\n",
        "saved_jsonl_path = f\"{PREPROCESSED_DIR}/{dataset_file}\"\n",
        "input_tensor_name = \"bytes_inputs\"  # Must match the serving_fn definition\n",
        "\n",
        "num_batches = -1  # uses the entire dataset\n",
        "start = time.time()\n",
        "all_train_images, all_train_labels = extract_images_and_labels(\n",
        "    train_ds, num_batches=num_batches\n",
        ")\n",
        "end = time.time()\n",
        "print(f\"Time taken to process training data: {end - start:.5f} secs\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ALNS_gFnTWjS"
      },
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "with open(saved_jsonl_path, \"w\") as f:\n",
        "    for i, im in enumerate(all_train_images):\n",
        "        img_bytes = io.BytesIO()\n",
        "        image = Image.fromarray(im.astype(np.uint8))\n",
        "        image.save(img_bytes, format=\"PNG\")\n",
        "        json.dump(\n",
        "            {\n",
        "                \"id\": str(i),\n",
        "                \"bytes_inputs\": {\n",
        "                    \"b64\": base64.b64encode(img_bytes.getvalue()).decode(\"utf-8\")\n",
        "                },\n",
        "            },\n",
        "            f,\n",
        "        )\n",
        "        f.write(\"\\n\")\n",
        "! gsutil -m cp {saved_jsonl_path} {BUCKET_URI}\n",
        "end = time.time()\n",
        "print(f\"Time taken to create and upload the training data: {end - start:.5f} secs\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "explanation_spec"
      },
      "source": [
        "### Example-based explanation specification\n",
        "\n",
        "Finally, you need to define the example-based explanation. To get explanations when doing a prediction, you must enable the explanation capability and set corresponding settings when you upload your custom model to an Vertex `Model` resource. These settings are referred to as the explanation metadata, which consists of:\n",
        "\n",
        "- `parameters`: This is the specification for the explainability algorithm to use for explanations on your model. In this tutorial, you will use `Examples`\n",
        "\n",
        "- `metadata`: This is the specification for how the algoithm is applied on your custom model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6DJ2FcXJc-bI"
      },
      "source": [
        "#### Explanation Parameters\n",
        "\n",
        "Let's first dive deeper into the settings for the explainability algorithm.\n",
        "\n",
        "#### Example-Based\n",
        "\n",
        "Example-based explanations enables analogy based explanations for data with applications in error analysis, model debugging, and batch labeling of new data. In turn, this can lead to more accurate and robust models, and efficient data labeling pipelines.\n",
        "\n",
        "\n",
        "Parameters:\n",
        "\n",
        "- `examples`: This paramenters allows to define conditions to return the nearest\n",
        "neighbors from the provided dataset.\n",
        "\n",
        "With Example-based explanations, you have a new explanation method with associated parameter configuration. Below you have the list of the main properties you have to define.\n",
        "\n",
        "- `dimensions` : The dimension of the embedding.\n",
        "- `approximateNeighborsCount` : Number of neighbors to return.\n",
        "- `distanceMeasureType` : The distance metric by which to measure nearness of examples. You can choose between ``SQUARED_L2_DISTANCE,  L1_DISTANCE, COSINE_DISTANCE and DOT_PRODUCT_DISTANCE``. \n",
        "- `featureNormType` : Normalize the embeddings so that it has a unit length. You can choose between ``UNIT_L2_NORM or NONE``.\n",
        "- `treeAhConfig`: Parameters controlling the trade-off between quality of approximation and speed. See the paper for technical details. Under the hood, it creates a shallow tree where the number of leaves is controlled by leafNodeEmbeddingCount and the search recall/speed tradeoff is controlled by leafNodesToSearchPercent.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sf_do58iXnK9"
      },
      "outputs": [],
      "source": [
        "DIMENSIONS = embedding_model.output.shape[1]\n",
        "DATASET_FILE_PATH = f\"{BUCKET_URI}/{dataset_file}\"\n",
        "\n",
        "NEAREST_NEIGHBOR_SEARCH_CONFIG = {\n",
        "    \"contentsDeltaUri\": \"\",\n",
        "    \"config\": {\n",
        "        \"dimensions\": DIMENSIONS,\n",
        "        \"approximateNeighborsCount\": 10,\n",
        "        \"distanceMeasureType\": \"SQUARED_L2_DISTANCE\",\n",
        "        \"featureNormType\": \"NONE\",\n",
        "        \"algorithmConfig\": {\n",
        "            \"treeAhConfig\": {\n",
        "                \"leafNodeEmbeddingCount\": 1000,\n",
        "                \"leafNodesToSearchPercent\": 100,\n",
        "            }\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "NUM_NEIGHBORS_TO_RETURN = 10\n",
        "\n",
        "EXAMPLES = vertex_ai_v1beta1.Examples(\n",
        "    nearest_neighbor_search_config=NEAREST_NEIGHBOR_SEARCH_CONFIG,\n",
        "    gcs_source=io_pb2.GcsSource(uris=[DATASET_FILE_PATH]),\n",
        "    neighbor_count=NUM_NEIGHBORS_TO_RETURN,\n",
        ")\n",
        "\n",
        "PARAMETERS = vertex_ai_v1beta1.ExplanationParameters(examples=EXAMPLES)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Dn_QepydOFv"
      },
      "source": [
        "#### Explanation Metadata\n",
        "\n",
        "Let's first dive deeper into the explanation metadata, which consists of:\n",
        "\n",
        "- `outputs`: It is represented by Map from output names to output metadata. In this case you expect embeddings.\n",
        "\n",
        "- `inputs`: It is represented by Metadata of the input of a feature. In this case you have the encoded image and the id associated to it. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tSMePRWgUaYs"
      },
      "outputs": [],
      "source": [
        "# for encoding parameter, 1 stands for 'IDENTITY'\n",
        "\n",
        "EXPLANATION_INPUTS = {\n",
        "    \"my_input\": vertex_ai_v1beta1.ExplanationMetadata.InputMetadata(\n",
        "        {\n",
        "            \"input_tensor_name\": input_tensor_name,\n",
        "            \"encoding\": vertex_ai_v1beta1.ExplanationMetadata.InputMetadata.Encoding(1),\n",
        "            \"modality\": \"image\",\n",
        "        }\n",
        "    ),\n",
        "    \"id\": vertex_ai_v1beta1.ExplanationMetadata.InputMetadata(\n",
        "        {\n",
        "            \"input_tensor_name\": \"id\",\n",
        "            \"encoding\": vertex_ai_v1beta1.ExplanationMetadata.InputMetadata.Encoding(1),\n",
        "        }\n",
        "    ),\n",
        "}\n",
        "\n",
        "EXPLANATION_OUTPUTS = {\n",
        "    \"embedding\": vertex_ai_v1beta1.ExplanationMetadata.OutputMetadata(\n",
        "        {\"output_tensor_name\": \"embedding\"}\n",
        "    )\n",
        "}\n",
        "\n",
        "EXPLANATION_META_CONFIG = vertex_ai_v1beta1.ExplanationMetadata(\n",
        "    inputs=EXPLANATION_INPUTS, outputs=EXPLANATION_OUTPUTS\n",
        ")\n",
        "\n",
        "EXPLANATION_SPEC = vertex_ai_v1beta1.ExplanationSpec(\n",
        "    parameters=PARAMETERS, metadata=EXPLANATION_META_CONFIG\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "upload_the_model:explanation"
      },
      "source": [
        "### Upload the model\n",
        "\n",
        "Use this helper function `upload_model` to upload your model, stored in SavedModel format, up to the `Model` service, which will instantiate a Vertex `Model` resource instance for your model. Once you've done that, you can use the `Model` resource instance in the same way as any other Vertex `Model` resource instance, such as deploying to an `Endpoint` resource for serving predictions.\n",
        "\n",
        "Let's now dive deeper into the Vertex model specification `model`. This is a dictionary object that consists of the following fields:\n",
        "\n",
        "- `display_name`: A human readable name for the `Model` resource.\n",
        "- `metadata_schema_uri`: Since your model was built without an Vertex `Dataset` resource, you will leave this blank (`''`).\n",
        "- `artificat_uri`: The Cloud Storage path where the embeddings is stored in SavedModel format.\n",
        "- `container_spec`: This is the specification for the Docker container that will be installed on the `Endpoint` resource, from which the `Model` resource will serve predictions. Use the variable you set earlier `DEPLOY_GPU != None` to use a GPU; otherwise only a CPU is allocated.\n",
        "- `explanation_spec`: This is the specification for enabling explainability for your model.\n",
        "\n",
        "The helper function those parameters in a configuration and calls the `Model` client service's method `upload_model`, which takes the following parameters:\n",
        "\n",
        "- `parent`: The Vertex location root path for `Dataset`, `Model` and `Endpoint` resources.\n",
        "- `model`: The specification for the Vertex `Model` resource instance.\n",
        "\n",
        "Uploading a model into a Vertex Model resource returns a long running operation, since it may take a few moments. You call response.result(), which is a synchronous call and will return when the Vertex Model resource is ready.\n",
        "\n",
        "The helper function returns the Vertex fully qualified identifier for the corresponding Vertex Model instance upload_model_response.model. You will save the identifier for subsequent steps in the variable model_to_deploy_id."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xATLh3ADhK-w"
      },
      "source": [
        "#### Define serving container configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z5ujqbVxhK-w"
      },
      "outputs": [],
      "source": [
        "DEPLOY_IMAGE_URI = \"gcr.io/cloud-aiplatform/prediction/tf2-cpu.2-5:latest\"\n",
        "\n",
        "CONTAINER_CONFIG = {\"image_uri\": DEPLOY_IMAGE_URI}\n",
        "\n",
        "CONTAINER_SPEC = vertex_ai_v1beta1.ModelContainerSpec(CONTAINER_CONFIG)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j0CR8R1KhTBV"
      },
      "source": [
        "#### Define Model configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7byz4C8wiRRj"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = f\"similarity-{DATASET_NAME}-{TIMESTAMP}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xh83zjREhTBW"
      },
      "outputs": [],
      "source": [
        "MODEL_CONFIGURATION = {\n",
        "    \"display_name\": MODEL_NAME,\n",
        "    \"artifact_uri\": EMBEDDINGS_URI,\n",
        "    \"metadata_schema_uri\": \"\",\n",
        "    \"container_spec\": CONTAINER_SPEC,\n",
        "    \"explanation_spec\": EXPLANATION_SPEC,\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b35eb7277ab4"
      },
      "source": [
        "This step can take up to an hour to finish. \n",
        "Currently, there is no easy way to monitor this progresses. Exposing the run logs and job status is planned for near future (in Preview)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kmFqtIO5OzvZ"
      },
      "outputs": [],
      "source": [
        "def upload_model(model_configuration):\n",
        "\n",
        "    model = vertex_ai_v1beta1.Model(\n",
        "        display_name=model_configuration[\"display_name\"],\n",
        "        artifact_uri=model_configuration[\"artifact_uri\"],\n",
        "        metadata_schema_uri=model_configuration[\"metadata_schema_uri\"],\n",
        "        explanation_spec=model_configuration[\"explanation_spec\"],\n",
        "        container_spec=model_configuration[\"container_spec\"],\n",
        "    )\n",
        "\n",
        "    response = clients[\"model\"].upload_model(parent=PARENT, model=model)\n",
        "    print(\"Long running operation:\", response.operation.name)\n",
        "    upload_model_response = response.result()\n",
        "    print(\"upload_model_response\")\n",
        "    print(\" model:\", upload_model_response.model)\n",
        "    return upload_model_response.model\n",
        "\n",
        "\n",
        "uploaded_model_id = upload_model(MODEL_CONFIGURATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "get_model"
      },
      "source": [
        "### Get `Model` resource information\n",
        "\n",
        "Now let's get the model information for just your model. Use this helper function `get_model`, with the following parameter:\n",
        "\n",
        "- `name`: The Vertex unique identifier for the `Model` resource.\n",
        "\n",
        "This helper function calls the Vertex `Model` client service's method `get_model`, with the following parameter:\n",
        "\n",
        "- `name`: The Vertex unique identifier for the `Model` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RtuXNfLdOzva"
      },
      "outputs": [],
      "source": [
        "def get_model(name):\n",
        "    response = clients[\"model\"].get_model(name=name)\n",
        "    print(response)\n",
        "\n",
        "\n",
        "get_model(uploaded_model_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_endpoint:custom"
      },
      "source": [
        "### Deploy the `Model` resource\n",
        "\n",
        "Now deploy the trained Vertex custom `Model` resource. This requires two steps:\n",
        "\n",
        "1. Create an `Endpoint` resource for deploying the `Model` resource to.\n",
        "\n",
        "2. Deploy the `Model` resource to the `Endpoint` resource."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_endpoint"
      },
      "source": [
        "#### Create an `Endpoint` resource\n",
        "\n",
        "Use this helper function `create_endpoint` to create an endpoint to deploy the model to for serving predictions, with the following parameter:\n",
        "\n",
        "- `display_name`: A human readable name for the `Endpoint` resource.\n",
        "\n",
        "The helper function uses the endpoint client service's `create_endpoint` method, which takes the following parameter:\n",
        "\n",
        "- `display_name`: A human readable name for the `Endpoint` resource.\n",
        "\n",
        "Creating an `Endpoint` resource returns a long running operation, since it may take a few moments to provision the `Endpoint` resource for serving. You call `response.result()`, which is a synchronous call and will return when the Endpoint resource is ready. The helper function returns the Vertex fully qualified identifier for the `Endpoint` resource: `response.name`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Qzc7H0jOzva"
      },
      "outputs": [],
      "source": [
        "ENDPOINT_NAME = f\"similarity-{DATASET_NAME}-endpoint-{TIMESTAMP}\"\n",
        "DESCRIPTION = \"An endpoint for the similarity model\"\n",
        "LABELS = {\"env\": ENVIRON, \"status\": \"online\"}\n",
        "\n",
        "\n",
        "def create_endpoint(display_name, description, labels):\n",
        "    endpoint = {\n",
        "        \"display_name\": display_name,\n",
        "        \"description\": description,\n",
        "        \"labels\": labels,\n",
        "    }\n",
        "    response = clients[\"endpoint\"].create_endpoint(parent=PARENT, endpoint=endpoint)\n",
        "    print(\"Long running operation:\", response.operation.name)\n",
        "\n",
        "    result = response.result()\n",
        "    print(\"result\")\n",
        "    print(\" name:\", result.name)\n",
        "    print(\" display_name:\", result.display_name)\n",
        "    print(\" description:\", result.description)\n",
        "    print(\" labels:\", result.labels)\n",
        "    print(\" create_time:\", result.create_time)\n",
        "    print(\" update_time:\", result.update_time)\n",
        "    return result\n",
        "\n",
        "\n",
        "result = create_endpoint(ENDPOINT_NAME, DESCRIPTION, LABELS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "endpoint_id:result"
      },
      "source": [
        "Now get the unique identifier for the `Endpoint` resource you created."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4pWgr2KkOzva"
      },
      "outputs": [],
      "source": [
        "# The full unique ID for the endpoint\n",
        "endpoint_id = result.name\n",
        "# The short numeric ID for the endpoint\n",
        "endpoint_short_id = endpoint_id.split(\"/\")[-1]\n",
        "\n",
        "print(endpoint_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "instance_scaling"
      },
      "source": [
        "#### Compute instance scaling\n",
        "\n",
        "You have several choices on scaling the compute instances for handling your online prediction requests:\n",
        "\n",
        "- Single Instance: The online prediction requests are processed on a single compute instance.\n",
        "  - Set the minimum (`MIN_NODES`) and maximum (`MAX_NODES`) number of compute instances to one.\n",
        "\n",
        "- Manual Scaling: The online prediction requests are split across a fixed number of compute instances that you manually specified.\n",
        "  - Set the minimum (`MIN_NODES`) and maximum (`MAX_NODES`) number of compute instances to the same number of nodes. When a model is first deployed to the instance, the fixed number of compute instances are provisioned and online prediction requests are evenly distributed across them.\n",
        "\n",
        "- Auto Scaling: The online prediction requests are split across a scaleable number of compute instances.\n",
        "  - Set the minimum (`MIN_NODES`) number of compute instances to provision when a model is first deployed and to de-provision, and set the maximum (`MAX_NODES`) number of compute instances to provision, depending on load conditions.\n",
        "\n",
        "The minimum number of compute instances corresponds to the field `min_replica_count` and the maximum number of compute instances corresponds to the field `max_replica_count`, in your subsequent deployment request."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8H20ToPEOzva"
      },
      "outputs": [],
      "source": [
        "MIN_NODES = 1\n",
        "MAX_NODES = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "deploy_model:dedicated,v1beta1"
      },
      "source": [
        "#### Deploy `Model` resource to the `Endpoint` resource\n",
        "\n",
        "Use this helper function `deploy_model` to deploy the model to the endpoint you created for serving predictions, with the following parameters:\n",
        "\n",
        "- `model`: The Vertex fully qualified identifier of the `Model` resource to upload (deploy) from the training pipeline.\n",
        "- `deploy_model_display_name`: A human readable name for the deployed model.\n",
        "- `endpoint`: The Vertex fully qualified `Endpoint` resource identifier to deploy the `Model` resource to.\n",
        "\n",
        "The helper function calls the `Endpoint` client service's method `deploy_model`, which takes the following parameters:\n",
        "\n",
        "- `endpoint`: The Vertex fully qualified `Endpoint` resource identifier to deploy the `Model` resource to to.\n",
        "- `deployed_model`: The requirements for deploying the model.\n",
        "- `traffic_split`: Percent of traffic at endpoint that goes to this model, which is specified as a dictioney of one or more key/value pairs.\n",
        "   - If only one model, then specify as **{ \"0\": 100 }**, where \"0\" refers to this model being uploaded and 100 means 100% of the traffic.\n",
        "   - If there are existing models on the endpoint, for which the traffic will be split, then specify as, where `model_id` is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
        "\n",
        "           { \"0\": percent, model_id: percent, ... }\n",
        "\n",
        "Let's now dive deeper into the `deployed_model` parameter. This parameter is specified as a Python dictionary with the minimum required fields:\n",
        "\n",
        "- `model`: The Vertex fully qualified identifier of the (upload) `Model` resource to deploy.\n",
        "- `display_name`: A human readable name for the deployed model.\n",
        "- `dedicated_resources`: This refers to how many compute instances (replicas) that are scaled for serving prediction requests.\n",
        "  - `machine_spec`: The compute instance to provision. Use the variable you set earlier `DEPLOY_GPU != None` to use a GPU; otherwise only a CPU is allocated.\n",
        "  - `min_replica_count`: The number of compute instances to initially provision, which you set earlier as the variable `MIN_NODES`.\n",
        "  - `max_replica_count`: The maximum number of compute instances to scale to, which you set earlier as the variable `MAX_NODES`.\n",
        "- `enable_container_logging`: This enables logging of container events, such as execution failures (default is container logging is disabled). Container logging is typically enabled when debugging the deployment and then disabled when deployed for production.\n",
        "\n",
        "#### Traffic Split\n",
        "\n",
        "Let's now dive deeper into the `traffic_split` parameter. This parameter is specified as a Python dictionary. This might at first be a tad bit confusing. Let me explain, you can deploy more than one instance of your model to an endpoint, and then set how much (percent) goes to each instance.\n",
        "\n",
        "Why would you do that? Perhaps you already have a previous version deployed in production -- let's call that v1. You got better model evaluation on v2, but you don't know for certain that it is really better until you deploy to production. So in the case of traffic split, you might want to deploy v2 to the same endpoint as v1, but it only get's say 10% of the traffic. That way, you can monitor how well it does without disrupting the majority of users -- until you make a final decision.\n",
        "\n",
        "#### Response\n",
        "\n",
        "The method returns a long running operation `response`. We will wait sychronously for the operation to complete by calling the `response.result()`, which will block until the model is deployed. If this is the first time a model is deployed to the endpoint, it may take a few additional minutes to complete provisioning of resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wtq7Xyn2Ozvb"
      },
      "outputs": [],
      "source": [
        "DEPLOYED_NAME = f\"similarity-{DATASET_NAME}-deployed-{TIMESTAMP}\"\n",
        "\n",
        "\n",
        "def deploy_model(\n",
        "    model, deployed_model_display_name, endpoint, traffic_split={\"0\": 100}\n",
        "):\n",
        "\n",
        "    if DEPLOY_GPU:\n",
        "        machine_spec = {\n",
        "            \"machine_type\": DEPLOY_COMPUTE,\n",
        "            \"accelerator_type\": DEPLOY_GPU,\n",
        "            \"accelerator_count\": DEPLOY_NGPU,\n",
        "        }\n",
        "    else:\n",
        "        machine_spec = {\n",
        "            \"machine_type\": DEPLOY_COMPUTE,\n",
        "            \"accelerator_count\": 0,\n",
        "        }\n",
        "\n",
        "    deployed_model = {\n",
        "        \"model\": model,\n",
        "        \"display_name\": deployed_model_display_name,\n",
        "        \"dedicated_resources\": {\n",
        "            \"min_replica_count\": MIN_NODES,\n",
        "            \"max_replica_count\": MAX_NODES,\n",
        "            \"machine_spec\": machine_spec,\n",
        "        },\n",
        "        \"enable_container_logging\": False,\n",
        "    }\n",
        "\n",
        "    response = clients[\"endpoint\"].deploy_model(\n",
        "        endpoint=endpoint, deployed_model=deployed_model, traffic_split=traffic_split\n",
        "    )\n",
        "\n",
        "    print(\"Long running operation:\", response.operation.name)\n",
        "    result = response.result()\n",
        "    print(\"result\")\n",
        "    deployed_model = result.deployed_model\n",
        "    print(\" deployed_model\")\n",
        "    print(\"  id:\", deployed_model.id)\n",
        "    print(\"  model:\", deployed_model.model)\n",
        "    print(\"  display_name:\", deployed_model.display_name)\n",
        "    print(\"  create_time:\", deployed_model.create_time)\n",
        "\n",
        "    return deployed_model.id\n",
        "\n",
        "\n",
        "deployed_model_id = deploy_model(uploaded_model_id, DEPLOYED_NAME, endpoint_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "make_prediction"
      },
      "source": [
        "## Query for similar examples\n",
        "\n",
        "Now do a online prediction to your deployed model to get your similar examples using a sample of validation dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAFuNUq5yr5s"
      },
      "source": [
        "### Prepare the validation data\n",
        "We will issue queries for this data. For the purpuse of the demo, we will choose a small subset of the full validation dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-BgUO1dEzNvU"
      },
      "outputs": [],
      "source": [
        "val_dataset_file = f\"{DATASET_NAME}-val-images.jsonl\"\n",
        "saved_val_jsonl_path = f\"{PREPROCESSED_DIR}/{val_dataset_file}\"\n",
        "\n",
        "num_batches = 10\n",
        "start = time.time()\n",
        "val_images, val_labels = extract_images_and_labels(\n",
        "    validation_ds, num_batches=num_batches\n",
        ")\n",
        "end = time.time()\n",
        "print(f\"Time taken to process validation data: {end - start:.5f} secs\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "prepare_test_item:test,image"
      },
      "source": [
        "### Prepare the request content\n",
        "You are going to send the STL10 image as compressed PNG image, instead of the raw uncompressed bytes. Also you encoded the bytes into base 64 -- which makes the content safe from modification when transmitting binary data over the network. You need to tell the serving binary where your model is deployed to, that the content has been base 64 encoded, so it will decode it on the other end in the serving binary.\n",
        "\n",
        "Each instance in the prediction request is a dictionary entry of the form:\n",
        "\n",
        "                        {`id`:, `bytes_inputs`: {'b64': content}}\n",
        "\n",
        "- `id`: the unique identifier associated to the image.\n",
        "- `bytes_inputs` : A map to contain decoded inputs.\n",
        "- `'b64'`: A key that indicates the content is base 64 encoded.\n",
        "- `content`: The compressed JPG image bytes as a base 64 encoded string."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wYoLoli7pL2v"
      },
      "outputs": [],
      "source": [
        "val_data = []\n",
        "\n",
        "for i, im in enumerate(val_images):\n",
        "    img_bytes = io.BytesIO()\n",
        "    image = Image.fromarray(im.astype(np.uint8))\n",
        "    image.save(img_bytes, format=\"PNG\")\n",
        "    instance = {\n",
        "        \"id\": str(i),\n",
        "        \"bytes_inputs\": {\"b64\": base64.b64encode(img_bytes.getvalue()).decode(\"utf-8\")},\n",
        "    }\n",
        "    val_data.append(instance)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "send_explain_request:image"
      },
      "source": [
        "### Send the prediction with explanation request\n",
        "\n",
        "Ok, now you have a test image. Use this helper function `explain_image`, which takes the parameters:\n",
        "\n",
        "- `image`: A list of test image data as a numpy array.\n",
        "- `endpoint`: The Vertex fully qualified identifier for the `Endpoint` resource where the `Model` resource was deployed.\n",
        "- `parameters_dict`: Additional parameters for serving.\n",
        "- `deployed_model_id`: The Vertex fully qualified identifier for the deployed model, when more than one model is deployed at the endpoint. Otherwise, if only one model deployed, can be set to `None`.\n",
        "\n",
        "This function uses the prediction client service and calls the `explain` method with the parameters:\n",
        "\n",
        "- `endpoint`: The Vertex fully qualified identifier for the `Endpoint` resource where the `Model` resource was deployed.\n",
        "- `instances`: A list of instances (encoded images) to predict and explain.\n",
        "- `parameters`: Additional parameters for serving.\n",
        "- `deployed_model_id`: The Vertex fully qualified identifier for the deployed model, when more than one model is deployed at the endpoint. Otherwise, if only one model deployed, can be set to `None`.\n",
        "\n",
        "Since the `predict()` service can take multiple images (instances), you will send your single image as a list of one image. As a final step, you package the instances list into Google's protobuf format -- which is what we pass to the `explain()` service.\n",
        "\n",
        "The `response` object returns a list, where each element in the list corresponds to the corresponding image in the request. You will see in the output for each prediction:\n",
        "\n",
        "- `deployed_model_id` -- The Vertex fully qualified identifer for the model that did the prediction/explanation.\n",
        "- `predictions` -- Confidence level for the prediction (`predictions`), between 0 and 1, for each of the ten classes.\n",
        "- `explanations` -- How each feature contributed to the prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "429YA19P14k7"
      },
      "outputs": [],
      "source": [
        "def explain_image(formatted_data, endpoint, parameters, deployed_model_id):\n",
        "\n",
        "    # The format of each instance should conform to the deployed model's prediction input schema.\n",
        "    instances_list = formatted_data\n",
        "    instances = [\n",
        "        json_format.ParseDict(instance, Value()) for instance in instances_list\n",
        "    ]\n",
        "\n",
        "    response = clients[\"prediction\"].explain(\n",
        "        endpoint=endpoint,\n",
        "        instances=instances,\n",
        "        parameters=parameters,\n",
        "        deployed_model_id=deployed_model_id,\n",
        "    )\n",
        "    print(\"response\")\n",
        "    print(\" deployed_model_id:\", response.deployed_model_id)\n",
        "    predictions = response.predictions\n",
        "    print(\"predictions\")\n",
        "    for prediction in predictions:\n",
        "        print(\" prediction:\", prediction)\n",
        "\n",
        "    explanations = response.explanations\n",
        "    print(\"explanations\")\n",
        "    for explanation in explanations:\n",
        "        print(\" explanation:\", explanation)\n",
        "\n",
        "    return response"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m76sOJ4803Qe"
      },
      "outputs": [],
      "source": [
        "# Update input and reformat to match the expected schema.\n",
        "BATCH_SIZE = (\n",
        "    8  # The request payload has a size limit so we need to subbatch our request\n",
        ")\n",
        "NUM_VAL_DATA = 35\n",
        "\n",
        "all_neighbors = []\n",
        "\n",
        "for data_idx in range(0, NUM_VAL_DATA, BATCH_SIZE):\n",
        "    end_idx = min(data_idx + BATCH_SIZE, NUM_VAL_DATA)\n",
        "    formatted_data = val_data[data_idx:end_idx]\n",
        "    response = explain_image(formatted_data, endpoint_id, None, deployed_model_id)\n",
        "    all_neighbors = (\n",
        "        all_neighbors + json_format.MessageToDict(response._pb)[\"explanations\"]\n",
        "    )\n",
        "\n",
        "print(f\"\\nExamples processed: {len(all_neighbors)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RZiJ1ykWIAWR"
      },
      "source": [
        "### Save input ids and the corresponding neighbors\n",
        "\n",
        "For each input image you sent, we create a dictionary with corrisponding neighbors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ILMgp_4LH9TO"
      },
      "outputs": [],
      "source": [
        "# Save input ids and the corresponding neighbors\n",
        "data_with_neighbors = []\n",
        "input_data_list = val_data[:NUM_VAL_DATA]\n",
        "\n",
        "for i, input_data in enumerate(input_data_list):\n",
        "    neighbor_dict = all_neighbors[i]\n",
        "    neighbor_dict[\"input\"] = input_data[\"id\"]\n",
        "    data_with_neighbors.append(neighbor_dict)\n",
        "\n",
        "DEBUG = False\n",
        "if DEBUG:\n",
        "    val_idx = 0\n",
        "    print(data_with_neighbors[val_idx])\n",
        "    print(data_with_neighbors[val_idx][\"neighbors\"])\n",
        "    print(data_with_neighbors[val_idx][\"input\"])\n",
        "    print(len(data_with_neighbors[val_idx][\"neighbors\"]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V9Gv-HxLlDOq"
      },
      "source": [
        "### Visualize the images with explanations\n",
        "\n",
        "In the following representation, you will see for each image sent the ten closer examples the API generated according the distance you define. \n",
        "\n",
        "As you can verify, although the `example index` results closed to image classified in the same category, in some cases the model wrongly indenfies the category. And you can easily visualize them by leveraging distances. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x79xz2XrxaNO"
      },
      "outputs": [],
      "source": [
        "val_img_indices = [1, 2, 10, 20, 25, 30]  # images to visually explore\n",
        "for val_img_idx in val_img_indices:\n",
        "    if val_img_idx > NUM_VAL_DATA - 1:\n",
        "        raise ValueError(\n",
        "            f\"Data index {val_img_idx} does not exist in the requested explanations\"\n",
        "        )\n",
        "    plot_input_and_neighbors(\n",
        "        val_img_idx,\n",
        "        all_train_images,\n",
        "        val_images,\n",
        "        all_train_labels,\n",
        "        val_labels,\n",
        "        label_index_to_name,\n",
        "        data_with_neighbors,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aLmS2M6dWsAI"
      },
      "source": [
        "## Further exploration\n",
        "If you want to continue exploring, here are some ideas:\n",
        "1.   Isolate test points where the model is making mistakes (cat mislabed as bird), and visualize the example-based explanations to see if you can find any common patterns.\n",
        "2.   If through this analysis, you find your training data is lacking in some representative cases (overhead images of cats), you can try adding such images to your dataset to see if that improves model performance.\n",
        "3.   [Fine-tune](https://keras.io/guides/transfer_learning/) the lower layers of the model to see if you can improve the quality of example-based explanations by enabling the model to learn a better latent representation.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "undeploy_model"
      },
      "source": [
        "## Undeploy the `Model` resource\n",
        "\n",
        "Now undeploy your `Model` resource from the serving `Endpoint` resoure. Use this helper function `undeploy_model`, which takes the following parameters:\n",
        "\n",
        "- `deployed_model_id`: The model deployment identifier returned by the endpoint service when the `Model` resource was deployed to.\n",
        "- `endpoint`: The Vertex fully qualified identifier for the `Endpoint` resource where the `Model` is deployed to.\n",
        "\n",
        "This function calls the endpoint client service's method `undeploy_model`, with the following parameters:\n",
        "\n",
        "- `deployed_model_id`: The model deployment identifier returned by the endpoint service when the `Model` resource was deployed.\n",
        "- `endpoint`: The Vertex fully qualified identifier for the `Endpoint` resource where the `Model` resource is deployed.\n",
        "- `traffic_split`: How to split traffic among the remaining deployed models on the `Endpoint` resource.\n",
        "\n",
        "Since this is the only deployed model on the `Endpoint` resource, you simply can leave `traffic_split` empty by setting it to {}."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RgKXH4qa4-5p"
      },
      "outputs": [],
      "source": [
        "def undeploy_model(deployed_model_id, endpoint):\n",
        "    response = clients[\"endpoint\"].undeploy_model(\n",
        "        endpoint=endpoint, deployed_model_id=deployed_model_id, traffic_split={}\n",
        "    )\n",
        "    print(response)\n",
        "\n",
        "\n",
        "undeploy_model(deployed_model_id, endpoint_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "To clean up all GCP resources used in this project, you can [delete the GCP\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
        "- Dataset\n",
        "- Pipeline\n",
        "- Model\n",
        "- Endpoint\n",
        "- Batch Job\n",
        "- Custom Job\n",
        "- Hyperparameter Tuning Job\n",
        "- Cloud Storage Bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6vBmkZzP4-5s"
      },
      "outputs": [],
      "source": [
        "delete_model = True\n",
        "delete_endpoint = True\n",
        "delete_customjob = True\n",
        "delete_bucket = True\n",
        "\n",
        "# Delete the model using the Vertex fully qualified identifier for the model\n",
        "try:\n",
        "    if delete_model and \"model_to_deploy_id\" in globals():\n",
        "        clients[\"model\"].delete_model(name=deployed_model_id)\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Delete the endpoint using the Vertex fully qualified identifier for the endpoint\n",
        "try:\n",
        "    if delete_endpoint and \"endpoint_id\" in globals():\n",
        "        clients[\"endpoint\"].delete_endpoint(name=endpoint_id)\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Delete the custom job using the Vertex fully qualified identifier for the custom job\n",
        "try:\n",
        "    if delete_customjob and \"job_id\" in globals():\n",
        "        clients[\"job\"].delete_custom_job(name=job_id)\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Delete bucket\n",
        "if delete_bucket and \"BUCKET_URI\" in globals():\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dc1622c1fd2f"
      },
      "outputs": [],
      "source": [
        "! rm custom.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YnT9ASjiocJN"
      },
      "outputs": [],
      "source": [
        "! rm -Rf {DATA_PATH} {DELIVERABLE_PATH} custom"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "showcase_custom_image_classification_online_explain_example_based_api.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}