community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46d6c164544f"
      },
      "outputs": [],
      "source": [
        "# Copyright 2020 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": "rEJBSTyZIrIb"
      },
      "source": [
        "# Training, Tuning and Deploying a PyTorch Text Classification Model on [Vertex AI](https://cloud.google.com/vertex-ai)\n",
        "## Fine-tuning pre-trained [BERT](https://huggingface.co/bert-base-cased) model for sentiment classification task "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X4cRE8IbIrIV"
      },
      "source": [
        "# Overview\n",
        "\n",
        "This example is inspired from Token-Classification [notebook](https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb) and [run_glue.py](https://github.com/huggingface/transformers/blob/v2.5.0/examples/run_glue.py). \n",
        "We will be fine-tuning **`bert-base-cased`** (pre-trained) model for sentiment classification task.\n",
        "You can find the details about this model at [Hugging Face Hub](https://huggingface.co/bert-base-cased).\n",
        "\n",
        "For more notebooks with the state of the art PyTorch/Tensorflow/JAX, you can explore [Hugging FaceNotebooks](https://huggingface.co/transformers/notebooks.html).\n",
        "\n",
        "### Dataset\n",
        "\n",
        "We will be using [IMDB movie review dataset](https://huggingface.co/datasets/imdb) from [Hugging Face Datasets](https://huggingface.co/datasets).\n",
        "\n",
        "### Objective\n",
        "\n",
        "How to **Build, Train, Tune and Deploy PyTorch models on [Vertex AI](https://cloud.google.com/vertex-ai)** and emphasize first class support for training and deploying PyTorch models on Vertex AI. \n",
        "\n",
        "### Table of Contents\n",
        "\n",
        "This notebook covers following sections:\n",
        "\n",
        "- [Creating Notebooks instance](#Creating-Notebooks-instance-on-Google-Cloud)\n",
        "- [Training](#Training)\n",
        "    - [Run Training Locally in the Notebook](#Training-locally-in-the-notebook)\n",
        "    - [Run Training Job on Vertex AI](#Training-on-Vertex-AI)\n",
        "        - [Training with pre-built container](#Run-Custom-Job-on-Vertex-AI-Training-with-a-pre-built-container)\n",
        "        - [Training with custom container](#Run-Custom-Job-on-Vertex-AI-Training-with-custom-container)\n",
        "- [Tuning](#Hyperparameter-Tuning) \n",
        "    - [Run Hyperparameter Tuning job on Vertex AI](#Run-Hyperparameter-Tuning-Job-on-Vertex-AI)\n",
        "- [Deploying](#Deploying)\n",
        "    - [Deploying model on Vertex AI Predictions with custom container](#Deploying-model-on-Vertex AI-Predictions-with-custom-container)\n",
        "\n",
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud Platform (GCP):\n",
        "\n",
        "* [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)\n",
        "* [Vertex AI Training](https://cloud.google.com/vertex-ai/docs/training/custom-training)\n",
        "* [Vertex AI Predictions](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions)\n",
        "* [Cloud Storage](https://cloud.google.com/storage)\n",
        "* [Container Registry](https://cloud.google.com/container-registry)\n",
        "* [Cloud Build](https://cloud.google.com/build) *[Optional]*\n",
        "\n",
        "Learn about [Vertex AI Pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage Pricing](https://cloud.google.com/storage/pricing) and [Cloud Build Pricing](https://cloud.google.com/build/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cda8ab8b7a04"
      },
      "source": [
        "## Creating Notebooks instance on Google Cloud\n",
        "\n",
        "This notebook assumes you are working with [PyTorch 1.9 DLVM](https://cloud.google.com/ai-platform/notebooks/docs/images) development environment with GPU runtime. You can create a Notebook instance using [Google Cloud Console](https://cloud.google.com/notebooks/docs/create-new) or [`gcloud` command](https://cloud.google.com/sdk/gcloud/reference/notebooks/instances/create).\n",
        "\n",
        "```\n",
        "gcloud notebooks instances create example-instance \\\n",
        "    --vm-image-project=deeplearning-platform-release \\\n",
        "    --vm-image-family=pytorch-1-9-cu110-notebooks \\\n",
        "    --machine-type=n1-standard-4 \\\n",
        "    --location=us-central1-a \\\n",
        "    --boot-disk-size=100 \\\n",
        "    --accelerator-core-count=1 \\\n",
        "    --accelerator-type=NVIDIA_TESLA_V100 \\\n",
        "    --install-gpu-driver \\\n",
        "    --network=default\n",
        "```\n",
        "***\n",
        "**NOTE:** You must have GPU quota before you can create instances with GPUs. Check the [quotas](https://console.cloud.google.com/iam-admin/quotas) page to ensure that you have enough GPUs available in your project. If GPUs are not listed on the quotas page or you require additional GPU quota, [request a quota increase](https://cloud.google.com/compute/quotas#requesting_additional_quota). Free Trial accounts do not receive GPU quota by default.\n",
        "***"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5178273783dd"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or Google Cloud Notebooks**, your environment already meets\n",
        "all the requirements to run this notebook. You can skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a1ead2cb5e3"
      },
      "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",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "1. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "1. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9a864c22307c"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Python dependencies required for this notebook are [Transformers](https://pypi.org/project/transformers/), [Datasets](https://pypi.org/project/datasets/) and [hypertune](https://github.com/GoogleCloudPlatform/cloudml-hypertune) will be installed in the Notebooks instance itself."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1fd00fa70a2a"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Google Cloud Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# Google Cloud Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2974254ea6be"
      },
      "outputs": [],
      "source": [
        "!pip -q install {USER_FLAG} --upgrade transformers\n",
        "!pip -q install {USER_FLAG} --upgrade datasets\n",
        "!pip -q install {USER_FLAG} --upgrade tqdm\n",
        "!pip -q install {USER_FLAG} --upgrade cloudml-hypertune"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0c1dcadc2c8"
      },
      "source": [
        "We will be using [Vertex AI SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#python) to interact with Vertex AI services. The high-level `aiplatform` library is designed to simplify common data science workflows by using wrapper classes and opinionated defaults. \n",
        "\n",
        "#### Install Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ab3fb1551abe"
      },
      "outputs": [],
      "source": [
        "!pip -q install {USER_FLAG} --upgrade google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f361541eff05"
      },
      "source": [
        "### Restart the Kernel\n",
        "\n",
        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2d77a223d63d"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import os\n",
        "\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": "9bebb3b46278"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Select a GPU runtime\n",
        "\n",
        "**Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select \"Runtime --> Change runtime type > GPU\"**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1e4d8f0c294"
      },
      "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",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "1. Enable following APIs in your project required for running the tutorial\n",
        "    - [Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)\n",
        "    - [Cloud Storage API](https://console.cloud.google.com/flows/enableapi?apiid=storage.googleapis.com)\n",
        "    - [Container Registry API](https://console.cloud.google.com/flows/enableapi?apiid=containerregistry.googleapis.com)\n",
        "    - [Cloud Build API](https://console.cloud.google.com/flows/enableapi?apiid=cloudbuild.googleapis.com)\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the 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": "36a4450b7c2e"
      },
      "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` or `google.auth`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "019e546007a3"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # <---CHANGE THIS TO YOUR PROJECT\n",
        "\n",
        "import os\n",
        "\n",
        "# Get your Google Cloud project ID using google.auth\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    import google.auth\n",
        "\n",
        "    _, PROJECT_ID = google.auth.default()\n",
        "    print(\"Project ID: \", PROJECT_ID)\n",
        "\n",
        "# validate PROJECT_ID\n",
        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
        "    print(\n",
        "        f\"Please set your project id before proceeding to next step. Currently it's set as {PROJECT_ID}\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c9906f72b18"
      },
      "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": "e90182316f63"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "\n",
        "def get_timestamp():\n",
        "    return datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "\n",
        "\n",
        "TIMESTAMP = get_timestamp()\n",
        "print(f\"TIMESTAMP = {TIMESTAMP}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4d96aec55eba"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "---\n",
        "\n",
        "**If you are using Google Cloud Notebooks**, your environment is already authenticated. Skip this step.\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82f63fb3401f"
      },
      "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** page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "2. Click **Create service account**.\n",
        "3. In the **Service account name** field, enter a name, and click **Create**.\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\" into the filter box, and select **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
        "5. Click *Create*. A JSON file that contains your key downloads to your local environment.\n",
        "6. Enter the path to your service account key as the `GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "535223fa4b84"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "# 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",
        "# The Google Cloud Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# If on Google Cloud Notebooks, then don't execute this code\n",
        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\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": "35094e21d888"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "When you submit a training job using the Cloud SDK, you upload a Python package containing your training code to a Cloud Storage bucket. Vertex AI runs the code from this package. In this tutorial, Vertex AI also saves the trained model that results from your job in the same bucket. Using this model artifact, you can then create Vertex AI model and endpoint resources in order to serve online predictions.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. It must be unique across all Cloud Storage buckets.\n",
        "\n",
        "You may also change the `REGION` variable, which is used for operations throughout the rest of this notebook. Make sure to [choose a region where Vertex AI services are available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions). You may not use a Multi-Regional Storage bucket for training with Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e07102312039"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # <---CHANGE THIS TO YOUR BUCKET\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "13b11a8299d6"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
        "    BUCKET_NAME = f\"gs://{PROJECT_ID}aip-{get_timestamp()}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4908f26b84be"
      },
      "outputs": [],
      "source": [
        "print(f\"PROJECT_ID = {PROJECT_ID}\")\n",
        "print(f\"BUCKET_NAME = {BUCKET_NAME}\")\n",
        "print(f\"REGION = {REGION}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6a3aae29644f"
      },
      "source": [
        "---\n",
        "\n",
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket.\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "25f9882bab87"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION $BUCKET_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2586e4ed72ad"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "315724257beb"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "276e063f41eb"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e894d41223e3"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import json\n",
        "import os\n",
        "import random\n",
        "import sys\n",
        "\n",
        "import google.auth\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform import gapic as aip\n",
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
        "from google.protobuf.json_format import MessageToDict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8330d87de404"
      },
      "outputs": [],
      "source": [
        "from IPython.display import HTML, display"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f656723ac8ed"
      },
      "outputs": [],
      "source": [
        "import datasets\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "import transformers\n",
        "from datasets import ClassLabel, Sequence, load_dataset\n",
        "from transformers import (AutoModelForSequenceClassification, AutoTokenizer,\n",
        "                          EvalPrediction, Trainer, TrainingArguments,\n",
        "                          default_data_collator)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29b773f66d1b"
      },
      "outputs": [],
      "source": [
        "print(f\"Notebook runtime: {'GPU' if torch.cuda.is_available() else 'CPU'}\")\n",
        "print(f\"PyTorch version : {torch.__version__}\")\n",
        "print(f\"Transformers version : {datasets.__version__}\")\n",
        "print(f\"Datasets version : {transformers.__version__}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9805f4ffaa4d"
      },
      "outputs": [],
      "source": [
        "APP_NAME = \"finetuned-bert-classifier\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4c7643572342"
      },
      "outputs": [],
      "source": [
        "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b3d8c26c6c98"
      },
      "source": [
        "# Training\n",
        "\n",
        "In this section, we will train a PyTorch model by fine-tuning pre-trained model from [Hugging Face Transformers](https://github.com/huggingface/transformers). We will train the model locally first and then on [Vertex AI training service](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ba110a456ac"
      },
      "source": [
        "## Training locally in the notebook"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whPRbBNbIrIl"
      },
      "source": [
        "### Loading the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7QYTpxXIrIl"
      },
      "source": [
        "For this example we will use [IMDB movie review dataset](https://huggingface.co/datasets/imdb) from [Hugging Face Datasets](https://huggingface.co/datasets/) for sentiment classification task. We use the [Hugging Face Datasets](https://github.com/huggingface/datasets) library to download the data. This can be easily done with the function `load_dataset`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IreSlFmlIrIm"
      },
      "outputs": [],
      "source": [
        "dataset = load_dataset(\"imdb\")\n",
        "dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RzfPtOMoIrIu"
      },
      "source": [
        "The `dataset` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GWiVUF0jIrIv"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"Total # of rows in training dataset {} and size {:5.2f} MB\".format(\n",
        "        dataset[\"train\"].shape[0], dataset[\"train\"].size_in_bytes / (1024 * 1024)\n",
        "    )\n",
        ")\n",
        "print(\n",
        "    \"Total # of rows in test dataset {} and size {:5.2f} MB\".format(\n",
        "        dataset[\"test\"].shape[0], dataset[\"test\"].size_in_bytes / (1024 * 1024)\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u3EtYfeHIrIz"
      },
      "source": [
        "To access an actual element, you need to select a split first, then give an index:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X6HrpprwIrIz"
      },
      "outputs": [],
      "source": [
        "dataset[\"train\"][0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dQPMGfEFLhFd"
      },
      "source": [
        "Using the `unique` method to extract label list. This will allow us to experiment with other datasets without hard-coding labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TYw3eIHfLhFd"
      },
      "outputs": [],
      "source": [
        "label_list = dataset[\"train\"].unique(\"label\")\n",
        "label_list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WHUmphG3IrI3"
      },
      "source": [
        "To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset (automatically decoding the labels in passing)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i3j8APAoIrI3"
      },
      "outputs": [],
      "source": [
        "def show_random_elements(dataset, num_examples=2):\n",
        "    assert num_examples <= len(\n",
        "        dataset\n",
        "    ), \"Can't pick more elements than there are in the dataset.\"\n",
        "    picks = []\n",
        "    for _ in range(num_examples):\n",
        "        pick = random.randint(0, len(dataset) - 1)\n",
        "        while pick in picks:\n",
        "            pick = random.randint(0, len(dataset) - 1)\n",
        "        picks.append(pick)\n",
        "\n",
        "    df = pd.DataFrame(dataset[picks])\n",
        "    for column, typ in dataset.features.items():\n",
        "        if isinstance(typ, ClassLabel):\n",
        "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
        "        elif isinstance(typ, Sequence) and isinstance(typ.feature, ClassLabel):\n",
        "            df[column] = df[column].transform(\n",
        "                lambda x: [typ.feature.names[i] for i in x]\n",
        "            )\n",
        "    display(HTML(df.to_html()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SZy5tRB_IrI7"
      },
      "outputs": [],
      "source": [
        "show_random_elements(dataset[\"train\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n9qywopnIrJH"
      },
      "source": [
        "### Preprocessing the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YVx71GdAIrJH"
      },
      "source": [
        "Before we can feed those texts to our model, we need to preprocess them. This is done by a pre-trained Hugging Face Transformers [`Tokenizer` class](https://huggingface.co/transformers/main_classes/tokenizer.html) which tokenizes the inputs (including converting the tokens to their corresponding IDs in the pre-trained vocabulary) and put it in a format the model expects, as well as generate the other inputs that model requires.\n",
        "\n",
        "To do all of this, we instantiate our tokenizer with the `AutoTokenizer.from_pretrained` method, which ensures:\n",
        "\n",
        "- we get a tokenizer that corresponds to the model architecture we want to use,\n",
        "- we download the vocabulary used when pre-training this specific checkpoint.\n",
        "\n",
        "That vocabulary will be cached, so it's not downloaded again the next time we run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "178b1743fd36"
      },
      "outputs": [],
      "source": [
        "batch_size = 16\n",
        "max_seq_length = 128\n",
        "model_name_or_path = \"bert-base-cased\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eXNLu_-nIrJI"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(\n",
        "    model_name_or_path,\n",
        "    use_fast=True,\n",
        ")\n",
        "# 'use_fast' ensure that we use fast tokenizers (backed by Rust) from the 🤗 Tokenizers library."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "muc7yghlLhFf"
      },
      "source": [
        "You can check type of models available with a fast tokenizer on the [big table of models](https://huggingface.co/transformers/index.html#bigtable)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rowT4iCLIrJK"
      },
      "source": [
        "You can directly call this tokenizer on one sentence:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a5hBlsrHIrJL"
      },
      "outputs": [],
      "source": [
        "tokenizer(\"Hello, this is one sentence!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qeT4rxOGLhFf"
      },
      "source": [
        "Depending on the model you selected, you will see different keys in the dictionary returned by the cell above. They don't matter much for what we're doing here (just know they are required by the model we will instantiate later), you can learn more about them in [this tutorial](https://huggingface.co/transformers/preprocessing.html) if you're interested.\n",
        "\n",
        "**NOTE:** If, as is the case here, your inputs have already been split into words, you should pass the list of words to your tokenizer with the argument `is_split_into_words=True`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MPunumj-LhFg"
      },
      "outputs": [],
      "source": [
        "example = dataset[\"train\"][4]\n",
        "print(example)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e3-m0sGHLhFg"
      },
      "outputs": [],
      "source": [
        "tokenizer(\n",
        "    [\"Hello\", \",\", \"this\", \"is\", \"one\", \"sentence\", \"split\", \"into\", \"words\", \".\"],\n",
        "    is_split_into_words=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G73hFICnLhFg"
      },
      "source": [
        "Note that transformers are often pre-trained with sub-word tokenizers, meaning that even if your inputs have been split into words already, each of those words could be split again by the tokenizer. Let's look at an example of that:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d8c6f3c11353"
      },
      "outputs": [],
      "source": [
        "# Dataset loading repeated here to make this cell idempotent\n",
        "# Since we are over-writing datasets variable\n",
        "dataset = load_dataset(\"imdb\")\n",
        "\n",
        "# Mapping labels to ids\n",
        "# NOTE: We can extract this automatically but the `Unique` method of the datasets\n",
        "# is not reporting the label -1 which shows up in the pre-processing.\n",
        "# Hence the additional -1 term in the dictionary\n",
        "label_to_id = {1: 1, 0: 0, -1: 0}\n",
        "\n",
        "\n",
        "def preprocess_function(examples):\n",
        "    \"\"\"\n",
        "    Tokenize the input example texts\n",
        "    NOTE: The same preprocessing step(s) will be applied\n",
        "    at the time of inference as well.\n",
        "    \"\"\"\n",
        "    args = (examples[\"text\"],)\n",
        "    result = tokenizer(\n",
        "        *args, padding=\"max_length\", max_length=max_seq_length, truncation=True\n",
        "    )\n",
        "\n",
        "    # Map labels to IDs (not necessary for GLUE tasks)\n",
        "    if label_to_id is not None and \"label\" in examples:\n",
        "        result[\"label\"] = [label_to_id[example] for example in examples[\"label\"]]\n",
        "\n",
        "    return result\n",
        "\n",
        "\n",
        "# apply preprocessing function to input examples\n",
        "dataset = dataset.map(preprocess_function, batched=True, load_from_cache_file=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "545PP3o8IrJV"
      },
      "source": [
        "### Fine-tuning the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FBiW8UpKIrJW"
      },
      "source": [
        "Now that our data is ready, we can download the pre-trained model and fine-tune it.\n",
        "\n",
        "---\n",
        "\n",
        "***Fine Tuning*** *involves taking a model that has already been trained for a given task and then tweaking the model for another similar task. Specifically, the tweaking involves replicating all the layers in the pre-trained model including weights and parameters, except the output layer. Then adding a new output classifier layer that predicts labels for the current task. The final step is to train the output layer from scratch, while the parameters of all layers from the pre-trained model are frozen. This allows learning from the pre-trained representations and \"fine-tuning\" the higher-order feature representations more relevant for the concrete task, such as analyzing sentiments in this case.* \n",
        "\n",
        "*For the scenario in the notebook analyzing sentiments, the pre-trained BERT model already encodes a lot of information about the language as the model was trained on a large corpus of English data in a self-supervised fashion. Now we only need to slightly tune them using their outputs as features for the sentiment classification task. This means quicker development iteration on a much smaller dataset, instead of training a specific Natural Language Processing (NLP) model with a larger training dataset.*\n",
        "\n",
        "\n",
        "![BERT fine-tuning](./images/bert-model-fine-tuning.png)\n",
        "\n",
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4d032e91d6c3"
      },
      "source": [
        "Since all our tasks are about token classification, we use the `AutoModelForSequenceClassification` class. Like with the tokenizer, the `from_pretrained` method will download and cache the model for us. The only thing we have to specify is the number of labels for our problem (which we can get from the features, as seen before):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TlqNaB8jIrJW"
      },
      "outputs": [],
      "source": [
        "model = AutoModelForSequenceClassification.from_pretrained(\n",
        "    model_name_or_path, num_labels=len(label_list)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CczA5lJlIrJX"
      },
      "source": [
        "**NOTE:** The warning is telling us we are throwing away some weights (the `vocab_transform` and `vocab_layer_norm` layers) and randomly initializing some other (the `pre_classifier` and `classifier` layers). This is absolutely normal in this case, because we are removing the head used to pre-train the model on a masked language modeling objective and replacing it with a new head for which we don't have pre-trained weights, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_N8urzhyIrJY"
      },
      "source": [
        "To instantiate a `Trainer`, we will need to define three more things. The most important is the [`TrainingArguments`](https://huggingface.co/transformers/main_classes/trainer.html#transformers.TrainingArguments), which is a class that contains all the attributes to customize the training. It requires one folder name, which will be used to save the checkpoints of the model, and all other arguments are optional:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bliy8zgjIrJY"
      },
      "outputs": [],
      "source": [
        "args = TrainingArguments(\n",
        "    evaluation_strategy=\"epoch\",\n",
        "    learning_rate=2e-5,\n",
        "    per_device_train_batch_size=batch_size,\n",
        "    per_device_eval_batch_size=batch_size,\n",
        "    num_train_epochs=1,\n",
        "    weight_decay=0.01,\n",
        "    output_dir=\"/tmp/cls\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "km3pGVdTIrJc"
      },
      "source": [
        "Here we set the evaluation to be done at the end of each epoch, tweak the learning rate, use the `batch_size` defined at the top of the notebook and customize the number of epochs for training, as well as the weight decay."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AqynGv2DLhFk"
      },
      "source": [
        "The last thing to define for our `Trainer` is how to compute the metrics from the predictions. You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a predictions and label_ids field) and has to return a dictionary string to float."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V233MNZgLhFk"
      },
      "outputs": [],
      "source": [
        "def compute_metrics(p: EvalPrediction):\n",
        "    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
        "    preds = np.argmax(preds, axis=1)\n",
        "    return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aehFUe14LhFl"
      },
      "source": [
        "Now we create the `Trainer` object and we are almost ready to train."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "imY1oC3SIrJf"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(\n",
        "    model,\n",
        "    args,\n",
        "    train_dataset=dataset[\"train\"],\n",
        "    eval_dataset=dataset[\"test\"],\n",
        "    data_collator=default_data_collator,\n",
        "    tokenizer=tokenizer,\n",
        "    compute_metrics=compute_metrics,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2b4255e2209f"
      },
      "source": [
        "You can add callbacks to the `trainer` object to customize the behavior of the training loop such as early stopping, reporting metrics at the end of evaluation phase or taking any decisions. In the hyperparameter tuning section of this notebook, we add a callback to `trainer` for automating hyperparameter tuning process."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CdzABDVcIrJg"
      },
      "source": [
        "We can now fine-tune our model by just calling the `train` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wZHmdWPELhFl"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4d3fcb41541c"
      },
      "outputs": [],
      "source": [
        "saved_model_local_path = \"./models\"\n",
        "!mkdir ./models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aac64fb35302"
      },
      "outputs": [],
      "source": [
        "trainer.save_model(saved_model_local_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKASz-2vIrJi"
      },
      "source": [
        "The `evaluate` method allows you to evaluate again on the evaluation dataset or on another dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UOUcBkX8IrJi"
      },
      "outputs": [],
      "source": [
        "history = trainer.evaluate()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d890a62c57c6"
      },
      "outputs": [],
      "source": [
        "history"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gX97-EH5LhFm"
      },
      "source": [
        "To get the other metrics computed  such as precision, recall or F1 score for each category, we can apply the same function as before on the result of the `predict` method."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cae4fdf26f2"
      },
      "source": [
        "### Run predictions locally with sample examples\n",
        "\n",
        "Using the trained model, we can predict the sentiment label for an input text after applying the preprocessing function that was used during the training. We will run the predictions locally in the notebook and later show how you can deploy the model to an endpoint using [TorchServe](https://pytorch.org/serve/) on Vertex AI Predictions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ab30d5518c"
      },
      "outputs": [],
      "source": [
        "model_name_or_path = \"bert-base-cased\"\n",
        "label_text = {0: \"Negative\", 1: \"Positive\"}\n",
        "saved_model_path = saved_model_local_path\n",
        "\n",
        "\n",
        "def predict(input_text, saved_model_path):\n",
        "    # initialize tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)\n",
        "\n",
        "    # preprocess and encode input text\n",
        "    tokenizer_args = (input_text,)\n",
        "    predict_input = tokenizer(\n",
        "        *tokenizer_args,\n",
        "        padding=\"max_length\",\n",
        "        max_length=128,\n",
        "        truncation=True,\n",
        "        return_tensors=\"pt\",\n",
        "    )\n",
        "\n",
        "    # load trained model\n",
        "    loaded_model = AutoModelForSequenceClassification.from_pretrained(saved_model_path)\n",
        "\n",
        "    # get predictions\n",
        "    output = loaded_model(predict_input[\"input_ids\"])\n",
        "\n",
        "    # return labels\n",
        "    label_id = torch.argmax(*output.to_tuple(), dim=1)\n",
        "\n",
        "    print(f\"Review text: {input_text}\")\n",
        "    print(f\"Sentiment : {label_text[label_id.item()]}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e594d72fa8e4"
      },
      "outputs": [],
      "source": [
        "# example #1\n",
        "review_text = (\n",
        "    \"\"\"Jaw dropping visual affects and action! One of the best I have seen to date.\"\"\"\n",
        ")\n",
        "predict_input = predict(review_text, saved_model_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dd393f0d38c1"
      },
      "outputs": [],
      "source": [
        "# example #2\n",
        "review_text = \"\"\"Take away the CGI and the A-list cast and you end up with film with less punch.\"\"\"\n",
        "predict_input = predict(review_text, saved_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "46a8d95366ee"
      },
      "source": [
        "## Training on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e97e8af2376c"
      },
      "source": [
        "You can do local experimentation on your Notebooks instance. However, for larger datasets or models often a vertically scaled compute or horizontally distributed training is required. The most effective way to perform this task is to leverage [Vertex AI custom training service](https://cloud.google.com/vertex-ai/docs/training/custom-training) for following reasons:\n",
        "\n",
        "- **Automatically provision and de-provision resources**: Training job on Vertex AI will automatically provision computing resources, performs the training task and ensures deletion of compute resources once the training job is finished.\n",
        "- **Reusability and portability**: You can package training code with its parameters and dependencies into a container and create a portable component. This container can then be run with different scenarios such as hyperparameter tuning, different data sources and more.\n",
        "- **Training at scale**: You can run a [distributed training job](https://cloud.google.com/vertex-ai/docs/training/distributed-training) with AI allowing you to train models in a cluster across multiple nodes in parallel and resulting in faster training time. \n",
        "- **Logging and Monitoring**: The training service logs messages from the job to [Cloud Logging](https://cloud.google.com/logging/docs) and can be monitored while the job is running.\n",
        "\n",
        "In this part of the notebook, we show how to scale the training job with Vertex AI by packaging the code and create a training pipeline to orchestrate a training job. There are three steps to run a training job using [Vertex AI custom training service](https://cloud.google.com/vertex-ai/docs/training/custom-training):\n",
        "\n",
        "- **STEP 1**: Determine training code structure - Packaging as a Python source distribution or as a custom container image\n",
        "- **STEP 2**: Chose a custom training method - custom job, hyperparameter training job or training pipeline\n",
        "- **STEP 3**: Run the training job\n",
        "\n",
        "![custom-training-on-vertex-ai](./images/custom-training-on-vertex-ai.png)\n",
        "\n",
        "#### Custom training methods\n",
        "\n",
        "There are three types of Vertex AI resources you can create to train custom models on Vertex AI:\n",
        "\n",
        "- **[Custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job):** With a custom job you configure the settings to run your training code on Vertex AI such as worker pool specs - machine types, accelerators, Python training spec or custom container spec. \n",
        "- **[Hyperparameter tuning jobs](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning):** Hyperparameter tuning jobs automate tuning of hyperparameters of your model based on the criteria you configure such as goal/metric to optimize, hyperparameters values and number of trials to run.\n",
        "- **[Training pipelines](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline):** Orchestrates custom training jobs or hyperparameter tuning jobs with additional steps after the training job is successfully completed.\n",
        "\n",
        "Please refer to the [documentation](https://cloud.google.com/vertex-ai/docs/training/custom-training-methods) for further details.\n",
        "\n",
        "In this notebook, we will cover Custom Jobs and Hyperparameter tuning jobs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "585ba542aa40"
      },
      "source": [
        "### Packaging the training application\n",
        "\n",
        "Before running the training job on Vertex AI, the training application code and any dependencies must be packaged and uploaded to Cloud Storage bucket or Container Registry or Artifact Registry that your Google Cloud project can access. This sections shows how to package and stage your application in the cloud.\n",
        "\n",
        "There are two ways to package your application and dependencies and train on Vertex AI:\n",
        "\n",
        "1. [Create a Python source distribution](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container) with the training code and dependencies to use with a [pre-built containers](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) on Vertex AI\n",
        "2. Use [custom containers](https://cloud.google.com/ai-platform/training/docs/custom-containers-training) to package dependencies using Docker containers\n",
        "\n",
        "**This notebook shows both packaging options to run a custom training job on Vertex AI.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b8b5333b577"
      },
      "source": [
        "#### Recommended Training Application Structure\n",
        "\n",
        "You can structure your training application in any way you like. However, the [following structure](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container#structure) is commonly used in Vertex AI samples, and having your project's organization be similar to the samples can make it easier for you to follow the samples.\n",
        "\n",
        "We have two directories `python_package` and `custom_container` showing both the packaging approaches. `README.md` files inside each directory has details on the directory structure and instructions on how to run application locally and on the cloud.\n",
        "\n",
        "```\n",
        ".\n",
        "├── custom_container\n",
        "│   ├── Dockerfile\n",
        "│   ├── README.md\n",
        "│   ├── scripts\n",
        "│   │   └── train-cloud.sh\n",
        "│   └── trainer -> ../python_package/trainer/\n",
        "├── python_package\n",
        "│   ├── README.md\n",
        "│   ├── scripts\n",
        "│   │   └── train-cloud.sh\n",
        "│   ├── setup.py\n",
        "│   └── trainer\n",
        "│       ├── __init__.py\n",
        "│       ├── experiment.py\n",
        "│       ├── metadata.py\n",
        "│       ├── model.py\n",
        "│       ├── task.py\n",
        "│       └── utils.py\n",
        "└── pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb    --> This notebook\n",
        "```\n",
        "\n",
        "1. Main project directory contains your `setup.py` file or `Dockerfile` with the dependencies. \n",
        "2. Use a subdirectory named `trainer` to store your main application module and `scripts` to submit training jobs locally or cloud\n",
        "3. Inside `trainer` directory:\n",
        "    - `task.py` - Main application module 1) initializes and parse task arguments (hyper parameters), and 2) entry point to the trainer\n",
        "    - `model.py` -  Includes function to create model with a sequence classification head from a pre-trained model.\n",
        "    - `experiment.py` - Runs the model training and evaluation experiment, and exports the final model.\n",
        "    - `metadata.py` - Defines metadata for classification task such as predefined model dataset name, target labels\n",
        "    - `utils.py` - Includes utility functions such as data input functions to read data, save model to GCS bucket"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7466d414a0e"
      },
      "source": [
        "### Run Custom Job on Vertex AI Training with a pre-built container"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a675e6ee9dc"
      },
      "source": [
        "Vertex AI provides Docker container images that can be run as [pre-built containers](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#available_container_images) for custom training. These containers include common dependencies used in training code based on the Machine Learning framework and framework version.\n",
        "\n",
        "In this notebook, we are using Hugging Face Datasets and fine tuning a transformer model from Hugging Face Transformers Library for sentiment analysis task using PyTorch. We will use [pre-built container for PyTorch](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#pytorch) and package the training application code by adding standard Python dependencies - `transformers`, `datasets` and `tqdm` - in the `setup.py` file. \n",
        "\n",
        "![Training with Prebuilt Containers on Vertex AI Training](./images/training-with-prebuilt-containers-on-vertex-training.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9481da82c2ca"
      },
      "source": [
        "Initialize the variables to define pre-built container image, location of training application and training module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "100eac26f547"
      },
      "outputs": [],
      "source": [
        "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest\"\n",
        ")\n",
        "\n",
        "PYTHON_PACKAGE_APPLICATION_DIR = \"python_package\"\n",
        "\n",
        "source_package_file_name = f\"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/trainer-0.1.tar.gz\"\n",
        "python_package_gcs_uri = (\n",
        "    f\"{BUCKET_NAME}/pytorch-on-gcp/{APP_NAME}/train/python_package/trainer-0.1.tar.gz\"\n",
        ")\n",
        "python_module_name = \"trainer.task\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2bdb8c21d0d5"
      },
      "source": [
        "Following is the `setup.py` file for the training application. The `find_packages()` function inside `setup.py` includes the `trainer` directory in the package as it contains `__init__.py` which tells [Python Setuptools](https://setuptools.readthedocs.io/en/latest/) to include all subdirectories of the parent directory as dependencies. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9dcb9b353ff3"
      },
      "outputs": [],
      "source": [
        "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/setup.py\n",
        "\n",
        "from setuptools import find_packages\n",
        "from setuptools import setup\n",
        "import setuptools\n",
        "\n",
        "from distutils.command.build import build as _build\n",
        "import subprocess\n",
        "\n",
        "\n",
        "REQUIRED_PACKAGES = [\n",
        "    'transformers',\n",
        "    'datasets',\n",
        "    'tqdm',\n",
        "    'cloudml-hypertune'\n",
        "]\n",
        "\n",
        "setup(\n",
        "    name='trainer',\n",
        "    version='0.1',\n",
        "    install_requires=REQUIRED_PACKAGES,\n",
        "    packages=find_packages(),\n",
        "    include_package_data=True,\n",
        "    description='Vertex AI | Training | PyTorch | Text Classification | Python Package'\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a90b5019436d"
      },
      "source": [
        "Run the following command to create a source distribution, dist/trainer-0.1.tar.gz:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9fe9d30e9806"
      },
      "outputs": [],
      "source": [
        "!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7b769681d231"
      },
      "source": [
        "Now upload the source distribution with training application to Cloud Storage bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a3492f5366aa"
      },
      "outputs": [],
      "source": [
        "!gsutil cp {source_package_file_name} {python_package_gcs_uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1e4a3c5a11a1"
      },
      "source": [
        "Validate the source distribution exists on Cloud Storage bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29fd8224d9cb"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -l {python_package_gcs_uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daf738249701"
      },
      "source": [
        "#### *[Optional]* **Run custom training job locally** "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dae35b9f526e"
      },
      "source": [
        "Before submitting the job to cloud, you can run the training job locally by calling the `trainer.task` module directly"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9defbf6fd074"
      },
      "outputs": [],
      "source": [
        "!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python -m trainer.task"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a1a4281ff7f"
      },
      "source": [
        "#### **Run custom training job on Vertex AI**\n",
        "\n",
        "We use [Vertex AI SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#client_libraries) to create and submit training job to the Vertex AI training service."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5d2957ef04fd"
      },
      "source": [
        "##### **Initialize the Vertex AI SDK for Python**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c4a578b55943"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6b0fed34b728"
      },
      "source": [
        "##### **Configure and submit Custom Job to Vertex AI Training service**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8f24d5bdc31d"
      },
      "source": [
        "Configure a [Custom Job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the [pre-built container](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) image for PyTorch and training code packaged as Python source distribution. \n",
        "\n",
        "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [Training Pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the Custom Job on Vertex AI Training service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "daababed95ac"
      },
      "outputs": [],
      "source": [
        "print(f\"APP_NAME={APP_NAME}\")\n",
        "print(\n",
        "    f\"PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI={PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI}\"\n",
        ")\n",
        "print(f\"python_package_gcs_uri={python_package_gcs_uri}\")\n",
        "print(f\"python_module_name={python_module_name}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1f7cff892eb8"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = f\"{APP_NAME}-pytorch-pkg-ar-{get_timestamp()}\"\n",
        "print(f\"JOB_NAME={JOB_NAME}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e7f68fec44ed"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=f\"{JOB_NAME}\",\n",
        "    python_package_gcs_uri=python_package_gcs_uri,\n",
        "    python_module_name=python_module_name,\n",
        "    container_uri=PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aaeb607c356e"
      },
      "outputs": [],
      "source": [
        "training_args = [\"--num-epochs\", \"2\", \"--model-name\", \"finetuned-bert-classifier\"]\n",
        "\n",
        "model = job.run(\n",
        "    replica_count=1,\n",
        "    machine_type=\"n1-standard-8\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "    accelerator_count=1,\n",
        "    args=training_args,\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bff8906cce7b"
      },
      "source": [
        "##### **Monitoring progress of the Custom Job**\n",
        "\n",
        "You can monitor the custom job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/training-pipelines/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)\n",
        "\n",
        "![Monitor custom job progress in Vertex AI Training](./images/vertex-training-monitor-custom-job.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "760ab4b88042"
      },
      "source": [
        "Validate the model artifacts written to GCS by the training code after the job completes successfully"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "39570bd8def9"
      },
      "outputs": [],
      "source": [
        "job_response = MessageToDict(job._gca_resource._pb)\n",
        "gcs_model_artifacts_uri = job_response[\"trainingTaskInputs\"][\"baseOutputDirectory\"][\n",
        "    \"outputUriPrefix\"\n",
        "]\n",
        "print(f\"Model artifacts are available at {gcs_model_artifacts_uri}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "09eaae731189"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -lr $gcs_model_artifacts_uri/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "813e2b1e6185"
      },
      "source": [
        "##### ***[Optional]*** **Submit custom job using gcloud CLI using Python source distribution**\n",
        "You can submit the training job to Vertex AI training service using [`gcloud beta ai custom-jobs create`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/create). `gcloud` command stages your training application on GCS bucket and submits the training job.\n",
        "\n",
        "\n",
        "```\n",
        "gcloud beta ai custom-jobs create \\\n",
        "    --display-name=${JOB_NAME} \\\n",
        "    --region ${REGION} \\\n",
        "    --python-package-uris=${PACKAGE_PATH} \\\n",
        "    --worker-pool-spec=replica-count=1,machine-type='n1-standard-8',accelerator-type='NVIDIA_TESLA_V100',accelerator-count=1,executor-image-uri=${IMAGE_URI},python-module='trainer.task',local-package-path=\"../python_package/\" \\\n",
        "    --args=\"--model-name\",\"finetuned-bert-classifier\",\"--job-dir\",$JOB_DIR\n",
        "```\n",
        "\n",
        "- `worker-pool-spec` parameter defines the worker pool configuration used by the custom job. Following are the fields within `worker-pool-spec`:\n",
        "    - Set the `executor-image-uri` to `us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest` for training on pre-built PyTorch v1.7 image for GPU\n",
        "    - Set the `local-package-path` to the path to the training code\n",
        "    - Set the `python-module` to the `trainer.task` which is the main module to start your application\n",
        "    - Set the `accelerator-type` and `machine-type` to set the compute type to run the application\n",
        " \n",
        "Refer [documentation](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/create#--args) for further details.\n",
        "\n",
        "The script at `./python_package/scripts/train-cloud.sh` contains the `gcloud` commands to launch the custom job and monitor the logs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d974c52416b8"
      },
      "outputs": [],
      "source": [
        "!cd python_package && ./scripts/train-cloud.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "edafdf2321d5"
      },
      "source": [
        "In addition to Cloud Console, you can monitor job progress by streaming logs using [gcloud CLI](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs) by passing the job id:\n",
        "\n",
        "\n",
        "```\n",
        "gcloud ai custom-jobs stream-logs <job_id> --region=$REGION\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06eb6c334000"
      },
      "source": [
        "You can validate the model artifacts written to GCS by the training code by running the following command:\n",
        "\n",
        "```\n",
        "!gsutil ls -l $JOB_DIR/\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c170d386492b"
      },
      "source": [
        "### Run Custom Job on Vertex AI Training with custom container"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "035227b6e581"
      },
      "source": [
        "To create a [training job with custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container?hl=hr), you define a `Dockerfile` to install or add the dependencies required for the training job. Then, you build and test your Docker image locally to verify, push the image to Container Registry and submit a Custom Job to Vertex AI Training service.\n",
        "\n",
        "![Training with custom containers on Vertex AI](./images/training-with-custom-containers-on-vertex-training.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95ea5367dc42"
      },
      "source": [
        "#### **Build your container using Dockerfile with Training Code and Dependencies**\n",
        "\n",
        "In the previous section, we wrapped the training application code and dependencies as Python source distribution. An alternate way to package the training application and dependencies is to [create a custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container?hl=hr) using Dockerfile. We create a Dockerfile with a [pre-built PyTorch container image provided by Vertex AI](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#available_container_images) as the base image, install the dependencies - `transformers`, `datasets` , `tqdm` and `cloudml-hypertune` and copy the training application code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "59cf98db5792"
      },
      "outputs": [],
      "source": [
        "%%writefile ./custom_container/Dockerfile\n",
        "\n",
        "# Use pytorch GPU base image\n",
        "FROM us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-10:latest\n",
        "\n",
        "# set working directory\n",
        "WORKDIR /app\n",
        "\n",
        "# Install required packages\n",
        "RUN pip install google-cloud-storage transformers datasets tqdm cloudml-hypertune\n",
        "\n",
        "# Copies the trainer code to the docker image.\n",
        "COPY ./trainer/__init__.py /app/trainer/__init__.py\n",
        "COPY ./trainer/experiment.py /app/trainer/experiment.py\n",
        "COPY ./trainer/utils.py /app/trainer/utils.py\n",
        "COPY ./trainer/metadata.py /app/trainer/metadata.py\n",
        "COPY ./trainer/model.py /app/trainer/model.py\n",
        "COPY ./trainer/task.py /app/trainer/task.py\n",
        "\n",
        "# Set up the entry point to invoke the trainer.\n",
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1f8d15e545e7"
      },
      "source": [
        "Build the image and tag the Container Registry path (gcr.io) that you will push to."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "03442268f05d"
      },
      "outputs": [],
      "source": [
        "CUSTOM_TRAIN_IMAGE_URI = f\"gcr.io/{PROJECT_ID}/pytorch_gpu_train_{APP_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "99d59a1f072c"
      },
      "outputs": [],
      "source": [
        "!cd ./custom_container/ && docker build -f Dockerfile -t $CUSTOM_TRAIN_IMAGE_URI ../python_package"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "169b86d50c45"
      },
      "source": [
        "#### *[Optional]* **Run Training Job Locally with Custom Container**\n",
        "\n",
        "Run the container locally in detached mode to test. When running with machine with GPUs, you can use `--gpus all` command line flag."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1361941a46d9"
      },
      "outputs": [],
      "source": [
        "!docker run --gpus all -it --rm $CUSTOM_TRAIN_IMAGE_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1873a991a40f"
      },
      "source": [
        "#### **Run custom training job on Vertex AI with Custom Container**\n",
        "\n",
        "Before submitting the training job to Vertex AI, push the custom container image to Google Cloud Container Registry and then submit the training job to Vertex AI. \n",
        "\n",
        "**NOTE:** *[Container Registry](https://cloud.google.com/container-registry) is a central repository to store, manage, and secure your Docker container images.*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d5e11e1319b6"
      },
      "source": [
        "##### **Push the container to Container Registry**\n",
        "\n",
        "Push your container image with training application code and dependencies to your Container Registry. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "11fbe35013f9"
      },
      "outputs": [],
      "source": [
        "!docker push $CUSTOM_TRAIN_IMAGE_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d157a91a792"
      },
      "source": [
        "Validate the custom container image in Container Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a23a4c447117"
      },
      "outputs": [],
      "source": [
        "!gcloud container images describe $CUSTOM_TRAIN_IMAGE_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a23e5e34bea9"
      },
      "source": [
        "##### **Initialize the Vertex AI SDK for Python**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "52f180952ca0"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "abf1fa4085cb"
      },
      "source": [
        "##### **Configure and submit Custom Job to Vertex AI Training service**\n",
        "\n",
        "Configure a [Custom Job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the [custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container) image with training code and other dependencies\n",
        "\n",
        "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [Training Pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the Custom Job to train on Vertex AI Training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5e048cc0500"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = f\"{APP_NAME}-pytorch-cstm-cntr-{get_timestamp()}\"\n",
        "\n",
        "print(f\"APP_NAME={APP_NAME}\")\n",
        "print(f\"CUSTOM_TRAIN_IMAGE_URI={CUSTOM_TRAIN_IMAGE_URI}\")\n",
        "print(f\"JOB_NAME={JOB_NAME}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6946d8ad77f6"
      },
      "outputs": [],
      "source": [
        "# configure the job with container image spec\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=f\"{JOB_NAME}\", container_uri=f\"{CUSTOM_TRAIN_IMAGE_URI}\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b7a24d15e8e9"
      },
      "outputs": [],
      "source": [
        "# define training code arguments\n",
        "training_args = [\"--num-epochs\", \"2\", \"--model-name\", \"finetuned-bert-classifier\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7459d3261450"
      },
      "outputs": [],
      "source": [
        "# submit the custom job to Vertex AI training service\n",
        "model = job.run(\n",
        "    replica_count=1,\n",
        "    machine_type=\"n1-standard-8\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "    accelerator_count=1,\n",
        "    args=training_args,\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d694283ad14c"
      },
      "source": [
        "##### **Monitoring progress of the Custom Job**\n",
        "\n",
        "You can monitor the custom job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/training-pipelines/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)\n",
        "\n",
        "![Monitor custom job progress in Vertex AI Training](./images/vertex-training-monitor-custom-job-container.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17d146293583"
      },
      "source": [
        "##### ***[Optional]*** **Submit Custom Job using gcloud CLI using custom container**\n",
        "You can submit the training job to Vertex AI training service using [`gcloud beta ai custom-jobs create`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/create) with custom container spec. `gcloud` command submits the training job and launches worker pool with the custom container image specified.\n",
        "\n",
        "\n",
        "```\n",
        "gcloud beta ai custom-jobs create \\\n",
        "    --display-name=${JOB_NAME} \\\n",
        "    --region ${REGION} \\\n",
        "    --worker-pool-spec=replica-count=1,machine-type='n1-standard-8',accelerator-type='NVIDIA_TESLA_V100',accelerator-count=1,container-image-uri=${CUSTOM_TRAIN_IMAGE_URI} \\\n",
        "    --args=\"--model-name\",\"finetuned-bert-classifier\",\"--job-dir\",$JOB_DIR\n",
        "```\n",
        "\n",
        "- `worker-pool-spec` parameter defines the worker pool configuration used by the custom job. Following are the fields within `worker-pool-spec`:\n",
        "    - Set the `container-image-uri` to the custom container image pushed to Google Cloud Container Registry for training\n",
        "    - Set the `accelerator-type` and `machine-type` to set the compute type to run the application\n",
        " \n",
        "Refer [documentation](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/create#--args) for further details.\n",
        "\n",
        "The script at `./custom_container/scripts/train-cloud.sh` contains the `gcloud` commands to launch the custom job and monitor the logs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "45c0f675d26a"
      },
      "outputs": [],
      "source": [
        "!cd custom_container && ./scripts/train-cloud.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69d723da4f80"
      },
      "source": [
        "In addition to Cloud Console, you can monitor job progress by streaming logs using [gcloud CLI](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs) by passing the job id:\n",
        "\n",
        "\n",
        "```\n",
        "gcloud ai custom-jobs stream-logs <job_id> --region=$REGION\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cffb0f2e5d4a"
      },
      "source": [
        "You can validate the model artifacts written to GCS by the training code by running the following command:\n",
        "\n",
        "```\n",
        "!gsutil ls -l $JOB_DIR/\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "409c7472fd4f"
      },
      "source": [
        "## Hyperparameter Tuning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ba6122f929e3"
      },
      "source": [
        "The training application code for fine-tuning a transformer model for sentiment analysis task uses hyperparameters such as learning rate and weight decay. These hyperparameters control the behavior of the training algorithm and can have a significant effect on the performance of the resulting model. This part of the notebook show how you can automate tuning these hyperparameters with Vertex AI Training service.\n",
        "\n",
        "We submit a [Hyperparameter Tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to Vertex AI Training service by packaging the training application code and dependencies in a Docker container and push the container to Google Container Registry, similar to running a Custom Job on Vertex AI with Custom Container.\n",
        "\n",
        "![Hyperparameter Tuning with Custom Containers on Vertex AI Training](./images/hp-tuning-with-custom-containers-on-vertex-training.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1570be63a7c"
      },
      "source": [
        "### How hyperparameter tuning works in Vertex AI?\n",
        "\n",
        "Following are the high level steps involved in running a Hyperparameter Tuning job on Vertex AI Training service:\n",
        "\n",
        "- You define the hyperparameters to tune the model along with the metric (or goal) to optimize\n",
        "- Vertex AI runs multiple trials of your training application with the hyperparameters and limits you specified - maximum number of trials to run and number of parallel trials. \n",
        "- Vertex AI keeps track of the results from each trial and makes adjustments for subsequent trials. This requires your training application to report the metrics to Vertex AI using the Python package [`cloudml-hypertune`](https://github.com/GoogleCloudPlatform/cloudml-hypertune). \n",
        "- When the job is finished, get the summary of all the trials with the most effective configuration of values based on the criteria you configured\n",
        "\n",
        "Refer to the [Vertex AI documentation](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to understand how to configure and select hyperparameters for tuning, configure tuning strategy and how Vertex AI optimizes the hyperparameter tuning jobs. The default tuning strategy uses results of previous trials to inform the assignment of values in subsequent trials. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d4dc37f4aab3"
      },
      "source": [
        "### Changes to training application code for hyperparameter tuning\n",
        "\n",
        "There are few [requirements](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview#the_flow_of_hyperparameter_values) to follow specific to hyperparameter tuning in Vertex AI:\n",
        "\n",
        "1. To pass the hyperparameter values to training code, you mist define a command-line argument in the main training module for each tuned hyperparameter. Use the value passed in those arguments to set the corresponding hyperparameter in the training application's code\n",
        "1. You must pass metrics from the training application to Vertex AI to evaluate the effectiveness of a trial. You can use [`cloudml-hypertune` Python package](https://github.com/GoogleCloudPlatform/cloudml-hypertune) to report metrics."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b0331ad2e11"
      },
      "source": [
        "Previously, in the training application code to fine-tune the transformer model for sentiment analysis task, we instantiated [`Trainer`](https://huggingface.co/transformers/main_classes/trainer.html) with hyperparameters passed as training arguments (`training_args`). \n",
        "\n",
        "```\n",
        "    # Estimator arguments\n",
        "    args_parser.add_argument(\n",
        "        '--learning-rate',\n",
        "        help='Learning rate value for the optimizers.',\n",
        "        default=2e-5,\n",
        "        type=float)\n",
        "    args_parser.add_argument(\n",
        "        '--weight-decay',\n",
        "        help=\"\"\"\n",
        "      The factor by which the learning rate should decay by the end of the\n",
        "      training.\n",
        "\n",
        "      decayed_learning_rate =\n",
        "        learning_rate * decay_rate ^ (global_step / decay_steps)\n",
        "\n",
        "      If set to 0 (default), then no decay will occur.\n",
        "      If set to 0.5, then the learning rate should reach 0.5 of its original\n",
        "          value at the end of the training.\n",
        "      Note that decay_steps is set to train_steps.\n",
        "      \"\"\",\n",
        "        default=0.01,\n",
        "        type=float)\n",
        "\n",
        "    # Enable hyperparameter\n",
        "    args_parser.add_argument(\n",
        "        '--hp-tune',\n",
        "        default=\"n\",\n",
        "        help='Enable hyperparameter tuning. Valida values are: \"y\" - enable, \"n\" - disable')\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "These hyperparameters are passed as command line arguments to the training module `trainer.task` which are then passed to the `training_args`. Refer to `./python_package/trainer` module for training application code.\n",
        "\n",
        "```\n",
        "    # set training arguments\n",
        "    training_args = TrainingArguments(\n",
        "        evaluation_strategy=\"epoch\",\n",
        "        learning_rate=args.learning_rate,\n",
        "        per_device_train_batch_size=args.batch_size,\n",
        "        per_device_eval_batch_size=args.batch_size,\n",
        "        num_train_epochs=args.num_epochs,\n",
        "        weight_decay=args.weight_decay,\n",
        "        output_dir=os.path.join(\"/tmp\", args.model_name)\n",
        "    )\n",
        "    \n",
        "    # initialize our Trainer\n",
        "    trainer = Trainer(\n",
        "        model,\n",
        "        training_args,\n",
        "        train_dataset=train_dataset,\n",
        "        eval_dataset=test_dataset,\n",
        "        data_collator=default_data_collator,\n",
        "        tokenizer=tokenizer,\n",
        "        compute_metrics=compute_metrics\n",
        "    )\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "16fb8c85a170"
      },
      "source": [
        "To report metrics to Vertex AI when hyperparameter tuning is enabled, we call [`cloudml-hypertune` Python package](https://github.com/GoogleCloudPlatform/cloudml-hypertune) after the evaluation phase which is added as a [callback](https://huggingface.co/transformers/main_classes/callback.html#transformers.trainer_callback.TrainerCallback) to the `trainer`. The `trainer` objects passes the metrics computed by the last evaluation phase to the callback which will be reported by `hypertune` library to Vertex AI for evaluating trials.\n",
        "\n",
        "```\n",
        "# add hyperparameter tuning callback to report metrics when enabled\n",
        "if args.hp_tune == \"y\":\n",
        "    trainer.add_callback(HPTuneCallback(\"accuracy\", \"eval_accuracy\"))\n",
        "\n",
        "class HPTuneCallback(TrainerCallback):\n",
        "    \"\"\"\n",
        "    A custom callback class that reports a metric to hypertuner\n",
        "    at the end of each epoch.\n",
        "    \"\"\"\n",
        "    \n",
        "    def __init__(self, metric_tag, metric_value):\n",
        "        super(HPTuneCallback, self).__init__()\n",
        "        self.metric_tag = metric_tag\n",
        "        self.metric_value = metric_value\n",
        "        self.hpt = hypertune.HyperTune()\n",
        "        \n",
        "    def on_evaluate(self, args, state, control, **kwargs):\n",
        "        print(f\"HP metric {self.metric_tag}={kwargs['metrics'][self.metric_value]}\")\n",
        "        self.hpt.report_hyperparameter_tuning_metric(\n",
        "            hyperparameter_metric_tag=self.metric_tag,\n",
        "            metric_value=kwargs['metrics'][self.metric_value],\n",
        "            global_step=state.epoch)\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a9d284f3599d"
      },
      "source": [
        "### Run Hyperparameter Tuning Job on Vertex AI\n",
        "\n",
        "Before submitting the hyperparameter tuning job to Vertex AI, push the custom container image with training application to Google Cloud Container Registry and then submit the job to Vertex AI. We will be using the same image used for running Custom Job on Vertex AI Training service."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0ab9c7321d2f"
      },
      "source": [
        "Validate the custom container image in Container Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "54c11399e7a9"
      },
      "outputs": [],
      "source": [
        "!gcloud container images describe $CUSTOM_TRAIN_IMAGE_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f60fab07d67c"
      },
      "source": [
        "##### **Initialize the Vertex AI SDK for Python**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7f914ea43ac0"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6652aa63ddff"
      },
      "source": [
        "##### **Configure and submit Hyperparameter Tuning Job to Vertex AI Training service**\n",
        "\n",
        "Configure a [Hyperparameter Tuning Job](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning) with the [custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container) image with training code and other dependencies.\n",
        "\n",
        "When configuring and submitting a Hyperparameter Tuning job, you need to attach a Custom Job definition with worker pool specs defining machine type, accelerators and URI for container image representing the custom container."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "93da68249c07"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = f\"{APP_NAME}-pytorch-hptune-{get_timestamp()}\"\n",
        "\n",
        "print(f\"APP_NAME={APP_NAME}\")\n",
        "print(f\"CUSTOM_TRAIN_IMAGE_URI={CUSTOM_TRAIN_IMAGE_URI}\")\n",
        "print(f\"JOB_NAME={JOB_NAME}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9d46db3a8b23"
      },
      "source": [
        "Define the training arguments with `hp-tune` argument set to `y` so that training application code can report metrics to Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "30b206352f6e"
      },
      "outputs": [],
      "source": [
        "training_args = [\n",
        "    \"--num-epochs\",\n",
        "    \"2\",\n",
        "    \"--model-name\",\n",
        "    \"finetuned-bert-classifier\",\n",
        "    \"--hp-tune\",\n",
        "    \"y\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0e84606af6ab"
      },
      "source": [
        "Create a **`CustomJob`** with worker pool specs to define machine types, accelerators and customer container spec with the training application code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0456de4efb46"
      },
      "outputs": [],
      "source": [
        "# The spec of the worker pools including machine type and Docker image\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": \"n1-standard-8\",\n",
        "            \"accelerator_type\": \"NVIDIA_TESLA_V100\",\n",
        "            \"accelerator_count\": 1,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\"image_uri\": CUSTOM_TRAIN_IMAGE_URI, \"args\": training_args},\n",
        "    }\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1339a33f3c3c"
      },
      "outputs": [],
      "source": [
        "custom_job = aiplatform.CustomJob(\n",
        "    display_name=JOB_NAME, worker_pool_specs=worker_pool_specs\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "269d6ffa999f"
      },
      "source": [
        "Define the `parameter_spec` as a Python dictionary object with the search space i.e. parameters to search and optimize. They key is the hyperparameter name passed as command line argument to the training code and value is the parameter specification. The spec requires to specify the hyperparameter data type as an instance of a parameter value specification. \n",
        "\n",
        "Refer to the [documentation](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview#hyperparameters) on selecting the hyperparaneter to tune and how to define parameter specification."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4210677c7ec7"
      },
      "outputs": [],
      "source": [
        "# Dictionary representing parameters to optimize.\n",
        "# The dictionary key is the parameter_id, which is passed into your training\n",
        "# job as a command line argument,\n",
        "# And the dictionary value is the parameter specification of the metric.\n",
        "parameter_spec = {\n",
        "    \"learning-rate\": hpt.DoubleParameterSpec(min=1e-6, max=0.001, scale=\"log\"),\n",
        "    \"weight-decay\": hpt.DiscreteParameterSpec(\n",
        "        values=[0.0001, 0.001, 0.01, 0.1], scale=None\n",
        "    ),\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e29058ada6d4"
      },
      "source": [
        "Define the `metric_spec` with name and goal of metric to optimize. The goal specifies whether you want to tune your model to maximize or minimize the value of this metric."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6d50f78b1303"
      },
      "outputs": [],
      "source": [
        "# Dictionary representing metrics to optimize.\n",
        "# The dictionary key is the metric_id, which is reported by your training job,\n",
        "# And the dictionary value is the optimization goal of the metric.\n",
        "metric_spec = {\"accuracy\": \"maximize\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3c5215ca590c"
      },
      "source": [
        "Configure and submit a Hyperparameter Tuning Job with the Custom Job, metric spec, parameter spec and trial limits.\n",
        "\n",
        "- **`max_trial_count`**: Maximum # of Trials run by the service. We recommend to start with a smaller value to understand the impact of the hyperparameters chosen before scaling up.\n",
        "- **`parallel_trial_count`**: Number of Trials to run in parallel. We recommend to start with a smaller value as Vertex AI uses results from the previous trials to inform the assignment of values in subsequent trials. Large # of parallel trials mean these trials start without having the benefit of the results of any trials still running.\n",
        "- **`search_algorithm`**: Search algorithm specified for the Study. If you do not specify an algorithm, Vertex AI by default applies Bayesian optimization to arrive at the optimal solution to search over the parameter space. \n",
        "\n",
        "Refer to the [documentation](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning#configuration) to understand the hyperparameter training job configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8d8088bfa3ad"
      },
      "outputs": [],
      "source": [
        "hp_job = aiplatform.HyperparameterTuningJob(\n",
        "    display_name=JOB_NAME,\n",
        "    custom_job=custom_job,\n",
        "    metric_spec=metric_spec,\n",
        "    parameter_spec=parameter_spec,\n",
        "    max_trial_count=5,\n",
        "    parallel_trial_count=2,\n",
        "    search_algorithm=None,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d4f07757c0a4"
      },
      "outputs": [],
      "source": [
        "model = hp_job.run(sync=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d07c7cb26ad8"
      },
      "source": [
        "##### **Monitoring progress of the Custom Job**\n",
        "\n",
        "You can monitor the hyperparameter tuning job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/hyperparameter-tuning-jobs/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)\n",
        "\n",
        "![Monitor hyperparameter tuning job progress in Vertex AI Training](./images/vertex-training-monitor-hptuning-job-container.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ba934b434f03"
      },
      "source": [
        "After the job is finished, you can view and format the results of the hyperparameter tuning Trials (run by Vertex AI Training service) as a Pandas dataframe"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9b5abf9ca2d7"
      },
      "outputs": [],
      "source": [
        "def get_trials_as_df(trials):\n",
        "    results = []\n",
        "    for trial in trials:\n",
        "        row = {}\n",
        "        t = MessageToDict(trial._pb)\n",
        "        # print(t)\n",
        "        row[\"Trial ID\"], row[\"Status\"], row[\"Start time\"], row[\"End time\"] = (\n",
        "            t[\"id\"],\n",
        "            t[\"state\"],\n",
        "            t[\"startTime\"],\n",
        "            t.get(\"endTime\", None),\n",
        "        )\n",
        "\n",
        "        for param in t[\"parameters\"]:\n",
        "            row[param[\"parameterId\"]] = param[\"value\"]\n",
        "\n",
        "        if t[\"state\"] == \"SUCCEEDED\":\n",
        "            row[\"Training step\"] = t[\"finalMeasurement\"][\"stepCount\"]\n",
        "            for metric in t[\"finalMeasurement\"][\"metrics\"]:\n",
        "                row[metric[\"metricId\"]] = metric[\"value\"]\n",
        "        results.append(row)\n",
        "\n",
        "    _df = pd.DataFrame(results)\n",
        "    return _df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "972f23712ce4"
      },
      "outputs": [],
      "source": [
        "df_trials = get_trials_as_df(hp_job.trials)\n",
        "df_trials"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5dbccb2b7d32"
      },
      "source": [
        "Now from the results of Trials, you can pick the best performing Trial to deploy to Vertex AI Predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c32f8a0a9c76"
      },
      "outputs": [],
      "source": [
        "# get trial id of the best run from the Trials\n",
        "best_trial_id = df_trials.loc[df_trials[\"accuracy\"].idxmax()][\"Trial ID\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8057eebb32e9"
      },
      "outputs": [],
      "source": [
        "# get base output directory where artifacts are saved\n",
        "base_output_dir = MessageToDict(hp_job._gca_resource._pb)[\"trialJobSpec\"][\n",
        "    \"baseOutputDirectory\"\n",
        "][\"outputUriPrefix\"]\n",
        "\n",
        "# get the model artifacts of the best trial id\n",
        "best_model_artifact_uri = f\"{base_output_dir}/{best_trial_id}\"\n",
        "\n",
        "print(\n",
        "    f\"Model artifacts from the Hyperparameter Tuning Job with bbest trial id {best_trial_id} are located at {best_model_artifact_uri}\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c120990dfa7e"
      },
      "source": [
        "You can validate the model artifacts written to GCS by the training code by running the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fbac4a0c28a5"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -r $best_model_artifact_uri/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2b43f9a01162"
      },
      "source": [
        "##### ***[Optional]*** **Submit hyperparameter tuning job using gcloud CLI**\n",
        "You can submit the hyperparameter tuning job to Vertex AI training service using [`gcloud beta ai hp-tuning-jobs create`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/hp-tuning-jobs/create). `gcloud` command submits the hyperparameter tuning job and launches multiple trials with worker pool based on custom container image specified and number of trials and the criteria set. The command requires hyperparameter tuning job configuration provided as configuration file in YAML format with job name.\n",
        "\n",
        "The following cell shows how to submit a hyperparameter tuning job on Vertex AI using `gcloud` CLI: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "39a793062138"
      },
      "outputs": [],
      "source": [
        "%%bash -s $BUCKET_NAME $APP_NAME\n",
        "\n",
        "# ========================================================\n",
        "# set job parameters\n",
        "# ========================================================\n",
        "# PROJECT_ID: Change to your project id\n",
        "PROJECT_ID=$(gcloud config list --format 'value(core.project)')\n",
        "        \n",
        "# set job display name\n",
        "JOB_PREFIX=\"finetuned-bert-classifier\"\n",
        "JOB_NAME=${JOB_PREFIX}-pytorch-hptune-$(date +%Y%m%d%H%M%S)\n",
        "echo \"Launching hyperparameter tuning job with display name as \"$JOB_NAME\n",
        "\n",
        "# BUCKET_NAME is a required parameter to run the cell.\n",
        "BUCKET_NAME=$1\n",
        "\n",
        "# APP_NAME: get application name\n",
        "APP_NAME=$2\n",
        "\n",
        "# JOB_DIR: Where to store prepared package and upload output model.\n",
        "JOB_DIR=${BUCKET_NAME}/${JOB_PREFIX}/model/${JOB_NAME}\n",
        "\n",
        "# custom container image URI\n",
        "CUSTOM_TRAIN_IMAGE_URI='gcr.io/'${PROJECT_ID}'/pytorch_gpu_train_'${APP_NAME}\n",
        "\n",
        "# ========================================================\n",
        "# create hyperparameter tuning configuration file\n",
        "# ========================================================\n",
        "cat << EOF > ./python_package/hptuning_job.yaml\n",
        "\n",
        "studySpec:\n",
        "  metrics:\n",
        "  - metricId: accuracy\n",
        "    goal: MAXIMIZE\n",
        "  parameters:\n",
        "  - parameterId: learning-rate\n",
        "    scaleType: UNIT_LOG_SCALE\n",
        "    doubleValueSpec:\n",
        "      minValue: 0.000001\n",
        "      maxValue: 0.001\n",
        "  - parameterId: weight-decay\n",
        "    scaleType: SCALE_TYPE_UNSPECIFIED\n",
        "    discreteValueSpec:\n",
        "      values: [\n",
        "          0.0001, \n",
        "          0.001, \n",
        "          0.01, \n",
        "          0.1\n",
        "      ]\n",
        "  measurementSelectionType: BEST_MEASUREMENT\n",
        "trialJobSpec:\n",
        "  workerPoolSpecs:\n",
        "  - machineSpec:\n",
        "      machineType: n1-standard-8\n",
        "      acceleratorType: NVIDIA_TESLA_V100\n",
        "      acceleratorCount: 1  \n",
        "    replicaCount: 1\n",
        "    containerSpec:\n",
        "      imageUri: $CUSTOM_TRAIN_IMAGE_URI\n",
        "      args: [\"--num-epochs\", \"2\", \"--model-name\", \"finetuned-bert-classifier\", \"--hp-tune\", \"y\"]\n",
        "  baseOutputDirectory: \n",
        "    outputUriPrefix: $JOB_DIR/\n",
        "EOF\n",
        "\n",
        "# ========================================================\n",
        "# submit hyperparameter tuning job\n",
        "# ========================================================\n",
        "gcloud beta ai hp-tuning-jobs create \\\n",
        "   --config ./python_package/hptuning_job.yaml \\\n",
        "   --display-name $JOB_NAME \\\n",
        "   --algorithm algorithm-unspecified \\\n",
        "   --max-trial-count 5 \\\n",
        "   --parallel-trial-count 2 \\\n",
        "   --region=us-central1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4685723ceaf"
      },
      "source": [
        "## Deploying\n",
        "\n",
        "Deploying a PyTorch model on [Vertex AI Predictions](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) requires to use a custom container that serves online predictions. You will deploy a container running [PyTorch's TorchServe](https://pytorch.org/serve/) tool in order to serve predictions from a fine-tuned transformer model from Hugging Face Transformers for sentiment analysis task. You can then use Vertex AI Predictions to classify sentiment of input texts. \n",
        "\n",
        "### Deploying model on Vertex AI Predictions with custom container\n",
        "\n",
        "To use a custom container to serve predictions from a PyTorch model, you must provide Vertex AI with a Docker container image that runs an HTTP server, such as TorchServe in this case. Please refer to [documentation](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) that describes the container image requirements to be compatible with Vertex AI Predictions.\n",
        "\n",
        "![Serving with Custom Containers on Vertex AI Predictions](./images/serve-pytorch-model-on-vertex-predictions-with-custom-containers.png)\n",
        "\n",
        "Essentially, to deploy a PyTorch model on Vertex AI Predictions following are the steps:\n",
        "\n",
        "1. Package the trained model artifacts including [default](https://pytorch.org/serve/#default-handlers) or [custom](https://pytorch.org/serve/custom_service.html) handlers by creating an archive file using [Torch model archiver](https://github.com/pytorch/serve/tree/master/model-archiver)\n",
        "2. Build a [custom container](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) compatible with Vertex AI Predictions to serve the model using Torchserve\n",
        "3. Upload the model with custom container image to serve predictions as a Vertex AI Model resource\n",
        "4. Create a Vertex AI Endpoint and [deploy the model](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) resource"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4c931da77e57"
      },
      "source": [
        "#### **Create a custom model handler to handle prediction requests**\n",
        "\n",
        "When predicting sentiments of the input text with the fine-tuned transformer model, it requires pre-processing of the input text and post-processing by adding name (positive/negative) to the target label (1/0) along with probability (or confidence). We create a custom handler script that is packaged with the model artifacts and TorchServe executes the code when it runs. \n",
        "\n",
        "Custom handler script does the following:\n",
        "\n",
        "- Pre-process input text before sending it to the model for inference\n",
        "- Customize how the model is invoked for inference\n",
        "- Post-process output from the model before sending back a response\n",
        "\n",
        "Please refer to the [TorchServe documentation](https://pytorch.org/serve/custom_service.html) for defining a custom handler."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ceb881740d6"
      },
      "outputs": [],
      "source": [
        "%%writefile predictor/custom_handler.py\n",
        "\n",
        "import os\n",
        "import json\n",
        "import logging\n",
        "\n",
        "import torch\n",
        "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
        "from ts.torch_handler.base_handler import BaseHandler\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "\n",
        "class TransformersClassifierHandler(BaseHandler):\n",
        "    \"\"\"\n",
        "    The handler takes an input string and returns the classification text \n",
        "    based on the serialized transformers checkpoint.\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super(TransformersClassifierHandler, self).__init__()\n",
        "        self.initialized = False\n",
        "\n",
        "    def initialize(self, ctx):\n",
        "        \"\"\" Loads the model.pt file and initialized the model object.\n",
        "        Instantiates Tokenizer for preprocessor to use\n",
        "        Loads labels to name mapping file for post-processing inference response\n",
        "        \"\"\"\n",
        "        self.manifest = ctx.manifest\n",
        "\n",
        "        properties = ctx.system_properties\n",
        "        model_dir = properties.get(\"model_dir\")\n",
        "        self.device = torch.device(\"cuda:\" + str(properties.get(\"gpu_id\")) if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "        # Read model serialize/pt file\n",
        "        serialized_file = self.manifest[\"model\"][\"serializedFile\"]\n",
        "        model_pt_path = os.path.join(model_dir, serialized_file)\n",
        "        if not os.path.isfile(model_pt_path):\n",
        "            raise RuntimeError(\"Missing the model.pt or pytorch_model.bin file\")\n",
        "        \n",
        "        # Load model\n",
        "        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)\n",
        "        self.model.to(self.device)\n",
        "        self.model.eval()\n",
        "        logger.debug('Transformer model from path {0} loaded successfully'.format(model_dir))\n",
        "        \n",
        "        # Ensure to use the same tokenizer used during training\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')\n",
        "\n",
        "        # Read the mapping file, index to object name\n",
        "        mapping_file_path = os.path.join(model_dir, \"index_to_name.json\")\n",
        "\n",
        "        if os.path.isfile(mapping_file_path):\n",
        "            with open(mapping_file_path) as f:\n",
        "                self.mapping = json.load(f)\n",
        "        else:\n",
        "            logger.warning('Missing the index_to_name.json file. Inference output will default.')\n",
        "            self.mapping = {\"0\": \"Negative\",  \"1\": \"Positive\"}\n",
        "\n",
        "        self.initialized = True\n",
        "\n",
        "    def preprocess(self, data):\n",
        "        \"\"\" Preprocessing input request by tokenizing\n",
        "            Extend with your own preprocessing steps as needed\n",
        "        \"\"\"\n",
        "        text = data[0].get(\"data\")\n",
        "        if text is None:\n",
        "            text = data[0].get(\"body\")\n",
        "        sentences = text.decode('utf-8')\n",
        "        logger.info(\"Received text: '%s'\", sentences)\n",
        "\n",
        "        # Tokenize the texts\n",
        "        tokenizer_args = ((sentences,))\n",
        "        inputs = self.tokenizer(*tokenizer_args,\n",
        "                                padding='max_length',\n",
        "                                max_length=128,\n",
        "                                truncation=True,\n",
        "                                return_tensors = \"pt\")\n",
        "        return inputs\n",
        "\n",
        "    def inference(self, inputs):\n",
        "        \"\"\" Predict the class of a text using a trained transformer model.\n",
        "        \"\"\"\n",
        "        prediction = self.model(inputs['input_ids'].to(self.device))[0].argmax().item()\n",
        "\n",
        "        if self.mapping:\n",
        "            prediction = self.mapping[str(prediction)]\n",
        "\n",
        "        logger.info(\"Model predicted: '%s'\", prediction)\n",
        "        return [prediction]\n",
        "\n",
        "    def postprocess(self, inference_output):\n",
        "        return inference_output\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f32fe0d49feb"
      },
      "source": [
        "##### **Generate target label to name file** *[Optional]*\n",
        "\n",
        "In the custom handler, we refer to a mapping file between target labels and their meaningful names that will be used to format the prediction response. Here we are mapping target label \"0\" as \"Negative\" and \"1\"  as \"Positive\". "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1be04c0f8960"
      },
      "outputs": [],
      "source": [
        "%%writefile ./predictor/index_to_name.json\n",
        "\n",
        "{\n",
        "    \"0\": \"Negative\", \n",
        "    \"1\": \"Positive\"\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d3feca8d40a3"
      },
      "source": [
        "#### **Create custom container image to serve predictions**\n",
        "\n",
        "We will use Cloud Build to create the custom container image with following build steps:\n",
        "\n",
        "##### **Download model artifacts**\n",
        "\n",
        "Download model artifacts that were saved as part of the training (or hyperparameter tuning) job from Cloud Storage to local directory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6966da1a55a5"
      },
      "outputs": [],
      "source": [
        "GCS_MODEL_ARTIFACTS_URI = best_model_artifact_uri"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "560170b3b7d9"
      },
      "source": [
        "Validate model artifact files in the Cloud Storage bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "216c2d276daa"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -r $GCS_MODEL_ARTIFACTS_URI/model/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a112fd084cf4"
      },
      "source": [
        "Copy files from Cloud Storage to local directory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "82e7e1a28c92"
      },
      "outputs": [],
      "source": [
        "!gsutil -m cp -r $GCS_MODEL_ARTIFACTS_URI/model/ ./predictor/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0302ecd8df0c"
      },
      "outputs": [],
      "source": [
        "!ls -ltrR ./predictor/model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0b44d55f7be1"
      },
      "source": [
        "##### **Build the container image**\n",
        "\n",
        "Create a Dockerfile with TorchServe as base image:\n",
        "\n",
        " - **`RUN`**: Installs dependencies such as `transformers`\n",
        " - **`COPY`**: Add model artifacts to `/home/model-server/` directory of the container image\n",
        " - **`COPY`**: Add custom handler script to `/home/model-server/` directory of the container image\n",
        " - **`RUN`**: Create `/home/model-server/config.properties` to define the serving configuration (health and prediction listener ports)\n",
        " - **`RUN`**: Run [Torch model archiver](https://pytorch.org/serve/model-archiver.html) to create a model archive file from the files copied into the image `/home/model-server/`. The model archive is saved in the `/home/model-server/model-store/` with name same as `<model-name>.mar`\n",
        " - **`CMD`**: Launch Torchserve HTTP server referencing the configuration properties and enables serving for the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "54acc94e5949"
      },
      "outputs": [],
      "source": [
        "%%bash -s $APP_NAME\n",
        "\n",
        "APP_NAME=$1\n",
        "\n",
        "cat << EOF > ./predictor/Dockerfile\n",
        "\n",
        "FROM pytorch/torchserve:latest-cpu\n",
        "\n",
        "# install dependencies\n",
        "RUN python3 -m pip install --upgrade pip\n",
        "RUN pip3 install transformers\n",
        "\n",
        "USER model-server\n",
        "\n",
        "# copy model artifacts, custom handler and other dependencies\n",
        "COPY ./custom_handler.py /home/model-server/\n",
        "COPY ./index_to_name.json /home/model-server/\n",
        "COPY ./model/$APP_NAME/ /home/model-server/\n",
        "\n",
        "# create torchserve configuration file\n",
        "USER root\n",
        "RUN printf \"\\nservice_envelope=json\" >> /home/model-server/config.properties\n",
        "RUN printf \"\\ninference_address=http://0.0.0.0:7080\" >> /home/model-server/config.properties\n",
        "RUN printf \"\\nmanagement_address=http://0.0.0.0:7081\" >> /home/model-server/config.properties\n",
        "USER model-server\n",
        "\n",
        "# expose health and prediction listener ports from the image\n",
        "EXPOSE 7080\n",
        "EXPOSE 7081\n",
        "\n",
        "# create model archive file packaging model artifacts and dependencies\n",
        "RUN torch-model-archiver -f \\\n",
        "  --model-name=$APP_NAME \\\n",
        "  --version=1.0 \\\n",
        "  --serialized-file=/home/model-server/pytorch_model.bin \\\n",
        "  --handler=/home/model-server/custom_handler.py \\\n",
        "  --extra-files \"/home/model-server/config.json,/home/model-server/tokenizer.json,/home/model-server/training_args.bin,/home/model-server/tokenizer_config.json,/home/model-server/special_tokens_map.json,/home/model-server/vocab.txt,/home/model-server/index_to_name.json\" \\\n",
        "  --export-path=/home/model-server/model-store\n",
        "\n",
        "# run Torchserve HTTP serve to respond to prediction requests\n",
        "CMD [\"torchserve\", \\\n",
        "     \"--start\", \\\n",
        "     \"--ts-config=/home/model-server/config.properties\", \\\n",
        "     \"--models\", \\\n",
        "     \"$APP_NAME=$APP_NAME.mar\", \\\n",
        "     \"--model-store\", \\\n",
        "     \"/home/model-server/model-store\"]\n",
        "EOF\n",
        "\n",
        "echo \"Writing ./predictor/Dockerfile\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e8878a9e37fe"
      },
      "source": [
        "Build the docker image tagged with Container Registry (gcr.io) path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d2d6b9a4857c"
      },
      "outputs": [],
      "source": [
        "CUSTOM_PREDICTOR_IMAGE_URI = f\"gcr.io/{PROJECT_ID}/pytorch_predict_{APP_NAME}\"\n",
        "print(f\"CUSTOM_PREDICTOR_IMAGE_URI = {CUSTOM_PREDICTOR_IMAGE_URI}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ac17dec3371"
      },
      "outputs": [],
      "source": [
        "!docker build \\\n",
        "  --tag=$CUSTOM_PREDICTOR_IMAGE_URI \\\n",
        "  ./predictor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5e13aeb92678"
      },
      "source": [
        "#### **Run the container locally** ***[Optional]***\n",
        "\n",
        "Before push the container image to Container Registry to use it with Vertex AI Predictions, you can run it as a container in your local environment to verify that the server works as expected"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eaf5bb210eb9"
      },
      "source": [
        "1. To run the container image as a container locally, run the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5b848c5fec2f"
      },
      "outputs": [],
      "source": [
        "!docker stop local_bert_classifier\n",
        "!docker run -t -d --rm -p 7080:7080 --name=local_bert_classifier $CUSTOM_PREDICTOR_IMAGE_URI\n",
        "!sleep 20"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a63999a79fea"
      },
      "source": [
        "2. To send the container's server a health check, run the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6499006af599"
      },
      "outputs": [],
      "source": [
        "!curl http://localhost:7080/ping"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a45d57cf5ba"
      },
      "source": [
        "If successful, the server returns the following response:\n",
        "\n",
        "```\n",
        "{\n",
        "  \"status\": \"Healthy\"\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6413202bfeb8"
      },
      "source": [
        "3. To send the container's server a prediction request, run the following commands:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "246d76732fcb"
      },
      "outputs": [],
      "source": [
        "%%bash -s $APP_NAME\n",
        "\n",
        "APP_NAME=$1\n",
        "\n",
        "cat > ./predictor/instances.json <<END\n",
        "{ \n",
        "   \"instances\": [\n",
        "     { \n",
        "       \"data\": {\n",
        "         \"b64\": \"$(echo 'Take away the CGI and the A-list cast and you end up with film with less punch.' | base64 --wrap=0)\"\n",
        "       }\n",
        "     }\n",
        "   ]\n",
        "}\n",
        "END\n",
        "\n",
        "curl -s -X POST \\\n",
        "  -H \"Content-Type: application/json; charset=utf-8\" \\\n",
        "  -d @./predictor/instances.json \\\n",
        "  http://localhost:7080/predictions/$APP_NAME/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d34a279f7fe7"
      },
      "source": [
        "This request uses a test sentence. If successful, the server returns prediction in below format:\n",
        "\n",
        "```\n",
        "    {\"predictions\": [\"Negative\"]}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ac697916268"
      },
      "source": [
        "4. To stop the container, run the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0a116f4fa88c"
      },
      "outputs": [],
      "source": [
        "!docker stop local_bert_classifier"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69477b3a00c0"
      },
      "source": [
        "#### **Deploying the serving container to Vertex AI Predictions**\n",
        "\n",
        "We create a model resource on Vertex AI and deploy the model to a Vertex AI Endpoints. You must deploy a model to an endpoint before using the model. The deployed model runs the custom container image to serve predictions. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "650995567bb1"
      },
      "source": [
        "##### **Push the serving container to Container Registry**\n",
        "\n",
        "Push your container image with inference code and dependencies to your Container Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "56a927f34d2a"
      },
      "outputs": [],
      "source": [
        "!docker push $CUSTOM_PREDICTOR_IMAGE_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a3da91e19af4"
      },
      "source": [
        "##### **Initialize the Vertex AI SDK for Python**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4e0857c467ef"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7a8e5f9ef00b"
      },
      "source": [
        "##### **Create a Model resource with custom serving container**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ad5f86202958"
      },
      "outputs": [],
      "source": [
        "VERSION = 1\n",
        "model_display_name = f\"{APP_NAME}-v{VERSION}\"\n",
        "model_description = \"PyTorch based text classifier with custom container\"\n",
        "\n",
        "MODEL_NAME = APP_NAME\n",
        "health_route = \"/ping\"\n",
        "predict_route = f\"/predictions/{MODEL_NAME}\"\n",
        "serving_container_ports = [7080]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c7fb34080960"
      },
      "outputs": [],
      "source": [
        "model = aiplatform.Model.upload(\n",
        "    display_name=model_display_name,\n",
        "    description=model_description,\n",
        "    serving_container_image_uri=CUSTOM_PREDICTOR_IMAGE_URI,\n",
        "    serving_container_predict_route=predict_route,\n",
        "    serving_container_health_route=health_route,\n",
        "    serving_container_ports=serving_container_ports,\n",
        ")\n",
        "\n",
        "model.wait()\n",
        "\n",
        "print(model.display_name)\n",
        "print(model.resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85134f5adfa5"
      },
      "source": [
        "For more context on upload or importing a model, refer [documentation](https://cloud.google.com/vertex-ai/docs/general/import-model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e12774f0dc68"
      },
      "source": [
        "##### **Create an Endpoint for Model with Custom Container**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d0bf59607969"
      },
      "outputs": [],
      "source": [
        "endpoint_display_name = f\"{APP_NAME}-endpoint\"\n",
        "endpoint = aiplatform.Endpoint.create(display_name=endpoint_display_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ba9826243edf"
      },
      "source": [
        "##### **Deploy the Model to Endpoint**\n",
        "\n",
        "Deploying a model associates physical resources with the model so it can serve online predictions with low latency. \n",
        "\n",
        "**NOTE:** This step takes few minutes to deploy the resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8e36ccbc24b4"
      },
      "outputs": [],
      "source": [
        "traffic_percentage = 100\n",
        "machine_type = \"n1-standard-4\"\n",
        "deployed_model_display_name = model_display_name\n",
        "min_replica_count = 1\n",
        "max_replica_count = 3\n",
        "sync = True\n",
        "\n",
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    deployed_model_display_name=deployed_model_display_name,\n",
        "    machine_type=machine_type,\n",
        "    traffic_percentage=traffic_percentage,\n",
        "    sync=sync,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bc4673478269"
      },
      "source": [
        "#### **Invoking the Endpoint with deployed Model using Vertex AI SDK to make predictions**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1fcbf1497407"
      },
      "source": [
        "##### **Get the Endpoint id**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "52362d052767"
      },
      "outputs": [],
      "source": [
        "endpoint_display_name = f\"{APP_NAME}-endpoint\"\n",
        "filter = f'display_name=\"{endpoint_display_name}\"'\n",
        "\n",
        "for endpoint_info in aiplatform.Endpoint.list(filter=filter):\n",
        "    print(\n",
        "        f\"Endpoint display name = {endpoint_info.display_name} resource id ={endpoint_info.resource_name} \"\n",
        "    )\n",
        "\n",
        "endpoint = aiplatform.Endpoint(endpoint_info.resource_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "376e74fdbd95"
      },
      "outputs": [],
      "source": [
        "endpoint.list_models()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1f454fff5e7d"
      },
      "source": [
        "##### **Formatting input for online prediction**\n",
        "\n",
        "This notebook uses [Torchserve's KServe based inference API](https://pytorch.org/serve/inference_api.html#kserve-inference-api) which is also [Vertex AI Predictions compatible format](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#prediction). For online prediction requests, format the prediction input instances as JSON with base64 encoding as shown here:\n",
        "\n",
        "```\n",
        "[\n",
        "    {\n",
        "        \"data\": {\n",
        "            \"b64\": \"<base64 encoded string>\"\n",
        "        }\n",
        "    }\n",
        "]\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "646e9dfe8add"
      },
      "source": [
        "Define sample texts to test predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "edc4658b2df1"
      },
      "outputs": [],
      "source": [
        "test_instances = [\n",
        "    b\"Jaw dropping visual affects and action! One of the best I have seen to date.\",\n",
        "    b\"Take away the CGI and the A-list cast and you end up with film with less punch.\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "52303da7c048"
      },
      "source": [
        "##### **Sending an online prediction request**\n",
        "\n",
        "Format input text string and call prediction endpoint with formatted input request and get the response"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a5808801dbcb"
      },
      "outputs": [],
      "source": [
        "print(\"=\" * 100)\n",
        "for instance in test_instances:\n",
        "    print(f\"Input text: \\n\\t{instance.decode('utf-8')}\\n\")\n",
        "    b64_encoded = base64.b64encode(instance)\n",
        "    test_instance = [{\"data\": {\"b64\": f\"{str(b64_encoded.decode('utf-8'))}\"}}]\n",
        "    print(f\"Formatted input: \\n{json.dumps(test_instance, indent=4)}\\n\")\n",
        "    prediction = endpoint.predict(instances=test_instance)\n",
        "    print(f\"Prediction response: \\n\\t{prediction}\")\n",
        "    print(\"=\" * 100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "21664e69720a"
      },
      "source": [
        "##### ***[Optional]*** **Make prediction requests using gcloud CLI**\n",
        "You can also call the Vertex AI Endpoint to make predictions using [`gcloud beta ai endpoints predict`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/endpoints/predict). \n",
        "\n",
        "The following cell shows how to make a prediction request to Vertex AI Endpoints using `gcloud` CLI: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7bb062d3e7a8"
      },
      "outputs": [],
      "source": [
        "endpoint_display_name = f\"{APP_NAME}-endpoint\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "950c04742e87"
      },
      "outputs": [],
      "source": [
        "%%bash -s $REGION $endpoint_display_name\n",
        "\n",
        "REGION=$1\n",
        "endpoint_display_name=$2\n",
        "\n",
        "# get endpoint id\n",
        "echo \"REGION = ${REGION}\"\n",
        "echo \"ENDPOINT DISPLAY NAME = ${endpoint_display_name}\"\n",
        "endpoint_id=$(gcloud beta ai endpoints list --region ${REGION} --filter \"display_name=${endpoint_display_name}\" --format \"value(ENDPOINT_ID)\")\n",
        "echo \"ENDPOINT_ID = ${endpoint_id}\"\n",
        "\n",
        "# call prediction endpoint\n",
        "input_text=\"Take away the CGI and the A-list cast and you end up with film with less punch.\"\n",
        "echo \"INPUT TEXT = ${input_text}\"\n",
        "\n",
        "prediction=$(\n",
        "echo \"\"\"\n",
        "{ \n",
        "   \"instances\": [\n",
        "     { \n",
        "       \"data\": {\n",
        "         \"b64\": \"$(echo ${input_text} | base64 --wrap=0)\"\n",
        "       }\n",
        "     }\n",
        "   ]\n",
        "}\n",
        "\"\"\" | gcloud beta ai endpoints predict ${endpoint_id} --region=$REGION --json-request -)\n",
        "\n",
        "echo \"PREDICTION RESPONSE = ${prediction}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e636b52e5913"
      },
      "source": [
        "## Cleaning up \n",
        "\n",
        "### Cleaning up training and deployment resources\n",
        "\n",
        "To clean up all Google Cloud resources used in this notebook, you can [delete the Google Cloud 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",
        "- Training Jobs\n",
        "- Model\n",
        "- Endpoint\n",
        "- Cloud Storage Bucket\n",
        "- Container Images\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7b89f5a348d1"
      },
      "source": [
        "Set flags for the resource type to be deleted"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "26ee41d02418"
      },
      "outputs": [],
      "source": [
        "delete_custom_job = False\n",
        "delete_hp_tuning_job = False\n",
        "delete_endpoint = True\n",
        "delete_model = False\n",
        "delete_bucket = False\n",
        "delete_image = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fa19a6540bd2"
      },
      "source": [
        "Define clients for jobs, models and endpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7f5f8dbd8cb6"
      },
      "outputs": [],
      "source": [
        "# API Endpoint\n",
        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
        "\n",
        "# Vertex AI location root path for your dataset, model and endpoint resources\n",
        "PARENT = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
        "\n",
        "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
        "\n",
        "# Initialize Vertex AI SDK\n",
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6e41b7b75ef6"
      },
      "outputs": [],
      "source": [
        "# functions to create client\n",
        "def create_job_client():\n",
        "    client = aip.JobServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "def create_model_client():\n",
        "    client = aip.ModelServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "def create_endpoint_client():\n",
        "    client = aip.EndpointServiceClient(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()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c41b12064678"
      },
      "source": [
        "Define functions to list the jobs, models and endpoints starting with `APP_NAME` defined earlier in the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c9a450d8eea2"
      },
      "outputs": [],
      "source": [
        "def list_custom_jobs():\n",
        "    client = clients[\"job\"]\n",
        "    jobs = []\n",
        "    response = client.list_custom_jobs(parent=PARENT)\n",
        "    for row in response:\n",
        "        _row = MessageToDict(row._pb)\n",
        "        if _row[\"displayName\"].startswith(APP_NAME):\n",
        "            jobs.append((_row[\"name\"], _row[\"displayName\"]))\n",
        "    return jobs\n",
        "\n",
        "\n",
        "def list_hp_tuning_jobs():\n",
        "    client = clients[\"job\"]\n",
        "    jobs = []\n",
        "    response = client.list_hyperparameter_tuning_jobs(parent=PARENT)\n",
        "    for row in response:\n",
        "        _row = MessageToDict(row._pb)\n",
        "        if _row[\"displayName\"].startswith(APP_NAME):\n",
        "            jobs.append((_row[\"name\"], _row[\"displayName\"]))\n",
        "    return jobs\n",
        "\n",
        "\n",
        "def list_models():\n",
        "    client = clients[\"model\"]\n",
        "    models = []\n",
        "    response = client.list_models(parent=PARENT)\n",
        "    for row in response:\n",
        "        _row = MessageToDict(row._pb)\n",
        "        if _row[\"displayName\"].startswith(APP_NAME):\n",
        "            models.append((_row[\"name\"], _row[\"displayName\"]))\n",
        "    return models\n",
        "\n",
        "\n",
        "def list_endpoints():\n",
        "    client = clients[\"endpoint\"]\n",
        "    endpoints = []\n",
        "    response = client.list_endpoints(parent=PARENT)\n",
        "    for row in response:\n",
        "        _row = MessageToDict(row._pb)\n",
        "        if _row[\"displayName\"].startswith(APP_NAME):\n",
        "            print(_row)\n",
        "            endpoints.append((_row[\"name\"], _row[\"displayName\"]))\n",
        "    return endpoints"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1281b5fe9222"
      },
      "source": [
        "#### **Deleting custom training jobs**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "944b5ec4694c"
      },
      "outputs": [],
      "source": [
        "# Delete the custom training using the Vertex AI fully qualified identifier for the custom training\n",
        "try:\n",
        "    if delete_custom_job:\n",
        "        custom_jobs = list_custom_jobs()\n",
        "        for job_id, job_name in custom_jobs:\n",
        "            print(f\"Deleting job {job_id} [{job_name}]\")\n",
        "            clients[\"job\"].delete_custom_job(name=job_id)\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd76a52356ec"
      },
      "source": [
        "#### **Deleting hyperparameter tuning jobs**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b65e57c07b43"
      },
      "outputs": [],
      "source": [
        "# Delete the hyperparameter tuning jobs using the Vertex AI fully qualified identifier for the hyperparameter tuning job\n",
        "try:\n",
        "    if delete_hp_tuning_job:\n",
        "        hp_tuning_jobs = list_hp_tuning_jobs()\n",
        "        for job_id, job_name in hp_tuning_jobs:\n",
        "            print(f\"Deleting job {job_id} [{job_name}]\")\n",
        "            clients[\"job\"].delete_hyperparameter_tuning_job(name=job_id)\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dc21008bb215"
      },
      "source": [
        "#### **Undeploy models and Delete endpoints**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "51684df7dc6b"
      },
      "outputs": [],
      "source": [
        "# Delete the endpoint using the Vertex AI fully qualified identifier for the endpoint\n",
        "try:\n",
        "    if delete_endpoint:\n",
        "        endpoints = list_endpoints()\n",
        "        for endpoint_id, endpoint_name in endpoints:\n",
        "            endpoint = aiplatform.Endpoint(endpoint_id)\n",
        "            # undeploy models from the endpoint\n",
        "            print(f\"Undeploying all deployed models from the endpoint {endpoint_name}\")\n",
        "            endpoint.undeploy_all(sync=True)\n",
        "            # deleting endpoint\n",
        "            print(f\"Deleting endpoint {endpoint_id} [{endpoint_name}]\")\n",
        "            clients[\"endpoint\"].delete_endpoint(name=endpoint_id)\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a6719f8f8a2"
      },
      "source": [
        "#### **Deleting models**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "087a5322f387"
      },
      "outputs": [],
      "source": [
        "# Delete the model using the Vertex AI fully qualified identifier for the model\n",
        "try:\n",
        "    models = list_models()\n",
        "    for model_id, model_name in models:\n",
        "        print(f\"Deleting model {model_id} [{model_name}]\")\n",
        "        clients[\"model\"].delete_model(name=model_id)\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "938a7cd172ab"
      },
      "source": [
        "#### **Delete contents from the staging bucket**\n",
        "\n",
        "---\n",
        "\n",
        "***NOTE: Everything in this Cloud Storage bucket will be DELETED. Please run it with caution.***\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a50bd867c02d"
      },
      "outputs": [],
      "source": [
        "if delete_bucket and \"BUCKET_NAME\" in globals():\n",
        "    print(f\"Deleting all contents from the bucket {BUCKET_NAME}\")\n",
        "\n",
        "    shell_output = ! gsutil du -as $BUCKET_NAME\n",
        "    print(\n",
        "        f\"Size of the bucket {BUCKET_NAME} before deleting = {shell_output[0].split()[0]} bytes\"\n",
        "    )\n",
        "\n",
        "    # uncomment below line to delete contents of the bucket\n",
        "    # ! gsutil rm -r $BUCKET_NAME\n",
        "\n",
        "    shell_output = ! gsutil du -as $BUCKET_NAME\n",
        "    if float(shell_output[0].split()[0]) > 0:\n",
        "        print(\n",
        "            \"PLEASE UNCOMMENT LINE TO DELETE BUCKET. CONTENT FROM THE BUCKET NOT DELETED\"\n",
        "        )\n",
        "\n",
        "    print(\n",
        "        f\"Size of the bucket {BUCKET_NAME} after deleting = {shell_output[0].split()[0]} bytes\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afd19ee7efa7"
      },
      "source": [
        "#### **Delete images from Container Registry**\n",
        "\n",
        "Deletes all the container images created in this tutorial with prefix defined by variable `APP_NAME` from the registry. All associated tags are also deleted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d45ec639148b"
      },
      "outputs": [],
      "source": [
        "gcr_images = !gcloud container images list --repository=gcr.io/$PROJECT_ID --filter=\"name~\"$APP_NAME\n",
        "\n",
        "if delete_image:\n",
        "    for image in gcr_images:\n",
        "        if image != \"NAME\":  # skip header line\n",
        "            print(f\"Deleting image {image} including all tags\")\n",
        "            !gcloud container images delete $image --force-delete-tags --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3541e1629c02"
      },
      "source": [
        "### Cleaning up Notebook Environment\n",
        "\n",
        "After you are done experimenting, you can either [STOP](https://cloud.google.com/ai-platform/notebooks/docs/shut-down) or DELETE the AI Notebook instance to prevent any  charges. If you want to save your work, you can choose to stop the instance instead.\n",
        "\n",
        "```\n",
        "# Stopping Notebook instance\n",
        "gcloud notebooks instances stop example-instance --location=us-central1-a\n",
        "\n",
        "\n",
        "# Deleting Notebook instance\n",
        "gcloud notebooks instances delete example-instance --location=us-central1-a\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "n9qywopnIrJH",
        "545PP3o8IrJV",
        "7k8ge1L1IrJk"
      ],
      "name": "pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}