notebooks/official/training/pytorch-text-sentiment-classification-custom-train-deploy.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "18ebbd838e32"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "16f0f41c2f9c"
      },
      "source": [
        "# Training, tuning and deploying a PyTorch text sentiment classification model on Vertex AI\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/pytorch-text-sentiment-classification-custom-train-deploy.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/pytorch-text-sentiment-classification-custom-train-deploy.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "<a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/training/pytorch-text-sentiment-classification-custom-train-deploy.ipynb\" target='_blank'>\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "     </a>\n",
        "  </td>\n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4339ebfdf7a0"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates building and deploying a text sentiment classification model by fine-tuing a pre-trained [BERT](https://huggingface.co/bert-base-cased) model using Vertex AI and Pytorch SDK. This example is inspired by the Hugging Face [Token_Classification](https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb) and [Run_Glue](https://github.com/huggingface/transformers/blob/v2.5.0/examples/run_glue.py) notebooks. \n",
        "\n",
        "You can find more details about the model at [Hugging Face Hub](https://huggingface.co/bert-base-cased). 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",
        "Learn more about [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9283a2954aef"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn to build, train, tune and deploy a PyTorch model on [Vertex AI](https://cloud.google.com/vertex-ai). You mainly focus on support for custom model training and deployment on Vertex AI. \n",
        "\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- Vertex AI `Workbench`\n",
        "- Vertex AI `Training`(Custom Python Package Training) \n",
        "- Vertex AI `Model Registry`\n",
        "- Vertex AI `Endpoint`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create training package for the text classification model.\n",
        "- Train the model with custom training on Vertex AI.\n",
        "- Check the created model artifacts.\n",
        "- Create a custom container for predictions.\n",
        "- Deploy the trained model to a Vertex AI Endpoint using the custom container for predictions.\n",
        "- Send online prediction requests to the deployed model and validate.\n",
        "- Clean up the resources created in this notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ab69c72f7c47"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [Happy Moments dataset](https://www.kaggle.com/ritresearch/happydb) from [Kaggle Datasets](https://www.kaggle.com/ritresearch/happydb). The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket.\n",
        "\n",
        "More information about this dataset can be found on [the HappyDB website](https://rit-public.github.io/HappyDB/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "181d4dfbf917"
      },
      "source": [
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "* Cloud Build\n",
        "* Artifact Registry\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), [Cloud Build pricing](https://cloud.google.com/build/pricing), [Artifact Registry pricing](https://cloud.google.com/artifact-registry/pricing) and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3848df1e5b0"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the packages required for executing this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6f967c29cbed"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Set the region\n",
        "\n",
        "**Optional**: Update the 'REGION' variable to specify the region that you want to use. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nsN5NJKSu-GU"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b45b2839f8b9"
      },
      "source": [
        "#### UUID\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 uuid for each instance session, and append the uuid onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e80050370d51"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import string\n",
        "\n",
        "\n",
        "# Generate a uuid of a specifed length(default=8)\n",
        "def generate_uuid(length: int = 8) -> str:\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
        "\n",
        "\n",
        "UUID = generate_uuid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "To authenticate your Google Cloud account, follow the instructions for your Jupyter environment:\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "<br>You are already authenticated.\n",
        "\n",
        "**2. Local JupyterLab instance**\n",
        "<br>Uncomment and run the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab**\n",
        "<br>Uncomment and run the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4d0f7629c309"
      },
      "source": [
        "### Import libraries and define constants\n",
        "\n",
        "Import the required libraries for this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "528bfbda0197"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import json\n",
        "import os\n",
        "import sys\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "from google.protobuf.json_format import MessageToDict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "290fd65fc20c"
      },
      "source": [
        "Define the constants needed for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9022ff0e121b"
      },
      "outputs": [],
      "source": [
        "# Name for the package application / model / repository\n",
        "APP_NAME = \"finetuned-bert-classifier\"\n",
        "\n",
        "# URI for the pre-built container for custom training\n",
        "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-11:latest\"\n",
        ")\n",
        "\n",
        "# Name of the folder where the python package needs to be stored\n",
        "PYTHON_PACKAGE_APPLICATION_DIR = \"python_package\"\n",
        "\n",
        "# Path to the source distribution tar of the python package\n",
        "source_package_file_name = f\"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/trainer-0.1.tar.gz\"\n",
        "\n",
        "# GCS path where the python package is stored\n",
        "python_package_gcs_uri = (\n",
        "    f\"{BUCKET_URI}/pytorch-on-gcp/{APP_NAME}/train/python_package/trainer-0.1.tar.gz\"\n",
        ")\n",
        "\n",
        "# Module name for training application\n",
        "python_module_name = \"trainer.task\"\n",
        "\n",
        "# Training job's display name\n",
        "JOB_NAME = f\"{APP_NAME}-pytorch-pkg-train-{UUID}\"\n",
        "\n",
        "# Set training job's machine-type\n",
        "TRAIN_MACHINE_TYPE = \"n1-standard-8\"\n",
        "# Set training job's accelerator type\n",
        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "# Set no. of h/w accelerators needed for the training job\n",
        "TRAIN_ACCELERATOR_COUNT = 1\n",
        "\n",
        "# Set the name of the container image for prediction\n",
        "CUSTOM_PREDICTOR_IMAGE_URI = (\n",
        "    f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{APP_NAME}/pytorch_predict_{APP_NAME}:latest\"\n",
        ")\n",
        "\n",
        "# Set the version for model-deployment\n",
        "VERSION = 1\n",
        "# Set the model display name\n",
        "model_display_name = f\"{APP_NAME}-v{VERSION}\"\n",
        "# Set the model description\n",
        "model_description = \"PyTorch based text classifier with custom container\"\n",
        "\n",
        "# Set the health route for prediction container\n",
        "health_route = \"/ping\"\n",
        "# Set the predict route for prediction container\n",
        "predict_route = f\"/predictions/{APP_NAME}\"\n",
        "# Set the serving container ports for prediction\n",
        "serving_container_ports = [7080]\n",
        "\n",
        "# Set the display name for endpoint\n",
        "endpoint_display_name = f\"{APP_NAME}-endpoint\"\n",
        "# Set the machine-type for deployment\n",
        "DEPLOY_MACHINE_TYPE = \"n1-standard-4\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d31e7ad1192"
      },
      "source": [
        "### Initialize the Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0d689f14d93c"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d52388cb4fd"
      },
      "source": [
        "## Custom Training on Vertex AI\n",
        "\n",
        "__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 organized similarly can make it easier for you to follow the samples.\n",
        "\n",
        "The following `python_package` directory structure shows a sample packaging approach.\n",
        "\n",
        "```\n",
        "├── python_package\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-sentiment-classification-custom-train-deploy.ipynb    --> This notebook\n",
        "```\n",
        "\n",
        "* Main project directory contains your `setup.py` file with the dependencies. \n",
        "* Inside `trainer` directory:\n",
        "    - `task.py` - Main application module initializes and parse task arguments (hyperparameters). It also serves as an entry point to the trainer.\n",
        "    - `model.py` -  Includes a function to create a 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 the metadata for classification tasks such as predefined model, dataset name and target labels.\n",
        "    - `utils.py` - Includes utility functions such as those used for reading data, saving models to Cloud Storage buckets."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ae8e41b828e1"
      },
      "source": [
        "### Create files requierd for the python package\n",
        "\n",
        "Create directories for the python package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d257b5c3c6b4"
      },
      "outputs": [],
      "source": [
        "!mkdir -p python_package/trainer\n",
        "!mkdir -p python_package/scripts\n",
        "!touch ./python_package/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "854420bdf10a"
      },
      "source": [
        "Create the `model.py` file that returns the specified pre-trained model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "112375025f5f"
      },
      "outputs": [],
      "source": [
        "%%writefile ./python_package/trainer/model.py\n",
        "\n",
        "from transformers import AutoModelForSequenceClassification\n",
        "from trainer import metadata\n",
        "\n",
        "def create(num_labels):\n",
        "    \"\"\"create the model by loading a pretrained model or define your \n",
        "    own\n",
        "\n",
        "    Args:\n",
        "      num_labels: number of target labels\n",
        "    \"\"\"\n",
        "    # Create the model, loss function, and optimizer\n",
        "    model = AutoModelForSequenceClassification.from_pretrained(\n",
        "        metadata.PRETRAINED_MODEL_NAME,\n",
        "        num_labels=num_labels\n",
        "    )\n",
        "    \n",
        "    return model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1fb1f3dfe245"
      },
      "source": [
        "Create the `utils.py` file that defines utility functions for data-loading, preprocessing, and model-saving."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d34c987cdb27"
      },
      "outputs": [],
      "source": [
        "%%writefile ./python_package/trainer/utils.py\n",
        "\n",
        "import os\n",
        "import datetime\n",
        "import pandas as pd\n",
        "\n",
        "from google.cloud import storage\n",
        "\n",
        "from transformers import AutoTokenizer\n",
        "from datasets import load_dataset, load_metric, ReadInstruction, DatasetDict, Dataset\n",
        "from trainer import metadata\n",
        "\n",
        "\n",
        "def preprocess_function(examples):\n",
        "    tokenizer = AutoTokenizer.from_pretrained(\n",
        "        metadata.PRETRAINED_MODEL_NAME,\n",
        "        use_fast=True,\n",
        "    )\n",
        "    \n",
        "    # Tokenize the texts\n",
        "    tokenizer_args = (\n",
        "        (examples['text'],) \n",
        "    )\n",
        "    result = tokenizer(*tokenizer_args, \n",
        "                       padding='max_length', \n",
        "                       max_length=metadata.MAX_SEQ_LENGTH, \n",
        "                       truncation=True)\n",
        "    \n",
        "    # We can extract this automatically but the unique() method of the dataset\n",
        "    # is not reporting the label -1 which shows up in the pre-processing\n",
        "    # hence the additional -1 term in the dictionary\n",
        "    \n",
        "    label_to_id = metadata.TARGET_LABELS\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[l] for l in examples[\"label\"]]\n",
        "\n",
        "    return result\n",
        "\n",
        "\n",
        "def load_data(args):\n",
        "    \"\"\"Loads the data into two different data loaders. (Train, Test)\n",
        "\n",
        "        Args:\n",
        "            args: arguments passed to the python script\n",
        "    \"\"\"\n",
        "    # dataset loading repeated here to make this cell idempotent\n",
        "    # since we are over-writing datasets variable\n",
        "    \n",
        "    df_train = pd.read_csv(metadata.TRAIN_DATA)\n",
        "    df_test = pd.read_csv(metadata.TEST_DATA)\n",
        "    \n",
        "    dataset = DatasetDict({\"train\": Dataset.from_pandas(df_train),\"test\": Dataset.from_pandas(df_test)})\n",
        "\n",
        "    dataset = dataset.map(preprocess_function, \n",
        "                          batched=True, \n",
        "                          load_from_cache_file=True)\n",
        "\n",
        "    train_dataset, test_dataset = dataset[\"train\"], dataset[\"test\"]\n",
        "\n",
        "    return train_dataset, test_dataset\n",
        "\n",
        "\n",
        "def save_model(args):\n",
        "    \"\"\"Saves the model to Google Cloud Storage or local file system\n",
        "\n",
        "    Args:\n",
        "      args: contains name for saved model.\n",
        "    \"\"\"\n",
        "    scheme = 'gs://'\n",
        "    if args.job_dir.startswith(scheme):\n",
        "        job_dir = args.job_dir.split(\"/\")\n",
        "        bucket_name = job_dir[2]\n",
        "        object_prefix = \"/\".join(job_dir[3:]).rstrip(\"/\")\n",
        "\n",
        "        if object_prefix:\n",
        "            model_path = '{}/{}'.format(object_prefix, args.model_name)\n",
        "        else:\n",
        "            model_path = '{}'.format(args.model_name)\n",
        "\n",
        "        bucket = storage.Client().bucket(bucket_name)    \n",
        "        local_path = os.path.join(\"/tmp\", args.model_name)\n",
        "        files = [f for f in os.listdir(local_path) if os.path.isfile(os.path.join(local_path, f))]\n",
        "        for file in files:\n",
        "            local_file = os.path.join(local_path, file)\n",
        "            blob = bucket.blob(\"/\".join([model_path, file]))\n",
        "            blob.upload_from_filename(local_file)\n",
        "        print(f\"Saved model files in gs://{bucket_name}/{model_path}\")\n",
        "    else:\n",
        "        print(f\"Saved model files at {os.path.join('/tmp', args.model_name)}\")\n",
        "        print(f\"To save model files in GCS bucket, please specify job_dir starting with gs://\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f048e99def15"
      },
      "source": [
        "Create the `metadata.py` file for defining the constants used in the training application."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "30715ff75292"
      },
      "outputs": [],
      "source": [
        "%%writefile ./python_package/trainer/metadata.py\n",
        "\n",
        "# Task type can be either 'classification', 'regression', or 'custom'.\n",
        "# This is based on the target feature in the dataset.\n",
        "TASK_TYPE = 'classification'\n",
        "\n",
        "# Dataset paths\n",
        "    \n",
        "TRAIN_DATA = \"gs://cloud-samples-data/ai-platform-unified/datasets/text/happydb/happydb_train.csv\"\n",
        "TEST_DATA = \"gs://cloud-samples-data/ai-platform-unified/datasets/text/happydb/happydb_test.csv\"\n",
        "\n",
        "# pre-trained model name\n",
        "PRETRAINED_MODEL_NAME = 'bert-base-cased'\n",
        "\n",
        "# List of the class values (labels) in a classification dataset.\n",
        "TARGET_LABELS = {\"leisure\": 0, \"exercise\":1, \"enjoy_the_moment\":2, \"affection\":3,\"achievement\":4, \"nature\":5, \"bonding\":6}\n",
        "\n",
        "\n",
        "# maximum sequence length\n",
        "MAX_SEQ_LENGTH = 128"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "26e7322f4639"
      },
      "source": [
        "Create the `experiment.py` file which defines the functions for hyperparameter tuning and training. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "049f50cd8671"
      },
      "outputs": [],
      "source": [
        "%%writefile ./python_package/trainer/experiment.py\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import hypertune\n",
        "\n",
        "from transformers import (\n",
        "    AutoTokenizer,\n",
        "    EvalPrediction,\n",
        "    Trainer,\n",
        "    TrainingArguments,\n",
        "    default_data_collator,\n",
        "    TrainerCallback\n",
        ")\n",
        "\n",
        "from trainer import model, metadata, utils\n",
        "\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",
        "\n",
        "\n",
        "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()}\n",
        "\n",
        "\n",
        "def train(args, model, train_dataset, test_dataset):\n",
        "    \"\"\"Create the training loop to load pretrained model and tokenizer and \n",
        "    start the training process\n",
        "\n",
        "    Args:\n",
        "      args: read arguments from the runner to set training hyperparameters\n",
        "      model: The neural network that you are training\n",
        "      train_dataset: The training dataset\n",
        "      test_dataset: The test dataset for evaluation\n",
        "    \"\"\"\n",
        "    \n",
        "    # initialize the tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(\n",
        "        metadata.PRETRAINED_MODEL_NAME,\n",
        "        use_fast=True,\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",
        "    \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",
        "    # training\n",
        "    trainer.train()\n",
        "    \n",
        "    return trainer\n",
        "\n",
        "\n",
        "def run(args):\n",
        "    \"\"\"Load the data, train, evaluate, and export the model for serving and\n",
        "     evaluating.\n",
        "\n",
        "    Args:\n",
        "      args: experiment parameters.\n",
        "    \"\"\"\n",
        "    # Open our dataset\n",
        "    train_dataset, test_dataset = utils.load_data(args)\n",
        "\n",
        "    label_list = train_dataset.unique(\"label\")\n",
        "    num_labels = len(label_list)\n",
        "    \n",
        "    # Create the model, loss function, and optimizer\n",
        "    text_classifier = model.create(num_labels=num_labels)\n",
        "    \n",
        "    # Train / Test the model\n",
        "    trainer = train(args, text_classifier, train_dataset, test_dataset)\n",
        "\n",
        "    metrics = trainer.evaluate(eval_dataset=test_dataset)\n",
        "    trainer.save_metrics(\"all\", metrics)\n",
        "\n",
        "    # Export the trained model\n",
        "    trainer.save_model(os.path.join(\"/tmp\", args.model_name))\n",
        "\n",
        "    # Save the model to GCS\n",
        "    if args.job_dir:\n",
        "        utils.save_model(args)\n",
        "    else:\n",
        "        print(f\"Saved model files at {os.path.join('/tmp', args.model_name)}\")\n",
        "        print(f\"To save model files in GCS bucket, please specify job_dir starting with gs://\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5aad46cadaaa"
      },
      "source": [
        "Create the `task.py` which is the main file that runs the training application."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ad78a00cba1"
      },
      "outputs": [],
      "source": [
        "%%writefile ./python_package/trainer/task.py\n",
        "\n",
        "import argparse\n",
        "import os\n",
        "\n",
        "from trainer import experiment\n",
        "\n",
        "\n",
        "def get_args():\n",
        "    \"\"\"Define the task arguments with the default values.\n",
        "\n",
        "    Returns:\n",
        "        experiment parameters\n",
        "    \"\"\"\n",
        "    args_parser = argparse.ArgumentParser()\n",
        "\n",
        "\n",
        "    # Experiment arguments\n",
        "    args_parser.add_argument(\n",
        "        '--batch-size',\n",
        "        help='Batch size for each training and evaluation step.',\n",
        "        type=int,\n",
        "        default=16)\n",
        "    args_parser.add_argument(\n",
        "        '--num-epochs',\n",
        "        help=\"\"\"\\\n",
        "        Maximum number of training data epochs on which to train.\n",
        "        If both --train-size and --num-epochs are specified,\n",
        "        --train-steps are: (train-size/train-batch-size) * num-epochs.\\\n",
        "        \"\"\",\n",
        "        default=1,\n",
        "        type=int,\n",
        "    )\n",
        "    args_parser.add_argument(\n",
        "        '--seed',\n",
        "        help='Random seed (default: 42)',\n",
        "        type=int,\n",
        "        default=42,\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 occurs.\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",
        "    # Saved model arguments\n",
        "    args_parser.add_argument(\n",
        "        '--job-dir',\n",
        "        default=os.getenv('AIP_MODEL_DIR'),\n",
        "        help='GCS location to export models')\n",
        "    args_parser.add_argument(\n",
        "        '--model-name',\n",
        "        default=\"finetuned-bert-classifier\",\n",
        "        help='The name of your saved model')\n",
        "\n",
        "    return args_parser.parse_args()\n",
        "\n",
        "\n",
        "def main():\n",
        "    \"\"\"Setup / Start the experiment\n",
        "    \"\"\"\n",
        "    args = get_args()\n",
        "    print(args)\n",
        "    experiment.run(args)\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92a191c72823"
      },
      "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 because it contains `__init__.py` which tells [Python Setuptools](https://setuptools.readthedocs.io/en/latest/) to include all subdirectories of the parent directory as dependencies. \n",
        "\n",
        "In `setup.py`, you also specify the Python packages that are required for the training application such as `transformers`, `datasets`, `cloudml-hypertune` and `tqdm`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1b4a2de28849"
      },
      "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": "cf62c83bdcbc"
      },
      "source": [
        "Run the following command to create a source distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "11d1299bcc12"
      },
      "outputs": [],
      "source": [
        "!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7e61169bcd7a"
      },
      "source": [
        "Now upload the source distribution with the training application to Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bff74cab1888"
      },
      "outputs": [],
      "source": [
        "!gsutil cp {source_package_file_name} {python_package_gcs_uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3d00f9beaba7"
      },
      "source": [
        "Validate that the source distribution exists in the Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "376bc46aa1e0"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -l {python_package_gcs_uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "88e4fd6394a0"
      },
      "source": [
        "### Run a custom job in Vertex AI using a pre-built container\n",
        "\n",
        "In this notebook, you are using Hugging Face Datasets and fine-tuning a transformer model from the Hugging Face Transformers library for sentiment analysis tasks using PyTorch. You don't need to build a PyTorch environment from scratch for running the training application because Vertex AI provides [pre-built containers](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#available_container_images). \n",
        "\n",
        "Vertex AI pre-built containers are Docker container images that you can use for custom training. They include some common dependencies used in training code based on the machine learning framework and framework version.\n",
        "\n",
        "You use a [pre-built container for PyTorch](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#pytorch) and the packaged training application to run the training job on Vertex AI.\n",
        "\n",
        "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. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c87d5618d366"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=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": "markdown",
      "metadata": {
        "id": "bfe709ea0a2f"
      },
      "source": [
        "Run the Custom training job with the following parameters:\n",
        "- `machine_type`: Mahcine type on which the job needs to run.\n",
        "- `accelerator_type`: Hardware accelerator type for running the job. One of _ACCELERATOR_TYPE_UNSPECIFIED_,\n",
        "        _NVIDIA_TESLA_K80_, _NVIDIA_TESLA_P100_, _NVIDIA_TESLA_V100_, _NVIDIA_TESLA_P4_,\n",
        "        _NVIDIA_TESLA_T4_, _NVIDIA_TELSA_A100_\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
        "- `replica_count`: The number of worker replicas.\n",
        "- `args`: Command line arguments to be passed to the Python script.\n",
        "\n",
        "Learn more about Vertex AI's [Custom Python-Package Trainining](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.CustomPythonPackageTrainingJob).\n",
        "\n",
        "*Note*: This training job may take over 24 hours."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b3709beb73e2"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING\"):\n",
        "    sys.exit(0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d8ac1bcb31a7"
      },
      "outputs": [],
      "source": [
        "training_args = [\"--num-epochs\", \"2\", \"--model-name\", APP_NAME]\n",
        "\n",
        "model = job.run(\n",
        "    replica_count=1,\n",
        "    machine_type=TRAIN_MACHINE_TYPE,\n",
        "    accelerator_type=TRAIN_ACCELERATOR_TYPE,\n",
        "    accelerator_count=TRAIN_ACCELERATOR_COUNT,\n",
        "    args=training_args,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d0f5ea96675"
      },
      "source": [
        "Validate that the model artifacts are written to Cloud Storage by the training code after the job completes successfully."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "502123b87fa3"
      },
      "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": "3d2205d5224f"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -lr $GCS_MODEL_ARTIFACTS_URI/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9caf6a802e99"
      },
      "source": [
        "## Deployment\n",
        "\n",
        "Deploying a PyTorch [model on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) requires you to use a custom container that serves online predictions on a Vertex AI Endpoint. You deploy a container running [PyTorch's TorchServe](https://pytorch.org/serve/) tool in order to serve predictions from the fine-tuned transformer model for a sentiment analysis task. Then, you can then use Vertex AI's online prediction service to classify the sentiment of input texts. \n",
        "\n",
        "### Deploying a model on Vertex AI using a 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 application, such as TorchServe in this case. Learn more about the [prediction container requirements on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements).\n",
        "\n",
        "Essentially, the following steps are needed to deploy a PyTorch model on Vertex AI:\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 to serve the model using TorchServe.\n",
        "3. Upload the model with the 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.\n",
        "\n",
        "\n",
        "#### Create a custom model handler to handle prediction requests\n",
        "\n",
        "When passing input text to the fine-tuned transformer model, the input text needs to be pre-processed. Once the model generates predictions, some post-processing has to be performed on the generated output to label it into the underlying classes and serve their probabilities (or confidence scores). \n",
        "\n",
        "To include the steps like pre-processing and post-processing, you create a custom handler script that is packaged with the model artifacts. Later, TorchServe executes the script when deployed. \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",
        "Learn more about defining a custom handler from [TorchServe documentation](https://pytorch.org/serve/custom_service.html).\n",
        "\n",
        "Create a directory to define a function for handling predictions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f109c2317225"
      },
      "outputs": [],
      "source": [
        "!mkdir -p predictor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9d6f534e5895"
      },
      "source": [
        "Create the `custom_handler.py` file that handles the prediction requests when deployed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e37fd93bcfa5"
      },
      "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 defaults.')\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": "2e504b5529b8"
      },
      "source": [
        "### Generate a file for class names\n",
        "\n",
        "For the custom handler, create the following mapping file (`index_to_name.json`) that is used to associate the target labels with their meaningful names while formatting the prediction responses."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d3c4969c2759"
      },
      "outputs": [],
      "source": [
        "%%writefile ./predictor/index_to_name.json\n",
        "\n",
        "{\n",
        "    \"0\": \"leisure\",\n",
        "    \"1\": \"exercise\",\n",
        "    \"2\": \"enjoy_the_moment\",\n",
        "    \"3\": \"affection\",\n",
        "    \"4\": \"achievement\",\n",
        "    \"5\": \"nature\",\n",
        "    \"6\": \"bonding\"\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6f0026a985fa"
      },
      "source": [
        "### Create a custom container image to serve predictions\n",
        "\n",
        "Next, you use [Artifact Registry](https://cloud.google.com/artifact-registry) and [Cloud Build](https://cloud.google.com/build) to create the custom container image in the following steps:\n",
        "\n",
        "#### Download the 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": "markdown",
      "metadata": {
        "id": "86b272c76889"
      },
      "source": [
        "Validate model artifact files in the Cloud Storage bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0d160b2911ac"
      },
      "outputs": [],
      "source": [
        "!gsutil ls -r $GCS_MODEL_ARTIFACTS_URI/model/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "38d171b3eb23"
      },
      "source": [
        "Copy the files from Cloud Storage to a local directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9138f330c9f2"
      },
      "outputs": [],
      "source": [
        "!gsutil -m cp -r $GCS_MODEL_ARTIFACTS_URI/model/ ./predictor/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2f9c5a82b1ed"
      },
      "outputs": [],
      "source": [
        "!ls -ltrR ./predictor/model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ec61feb1188"
      },
      "source": [
        "#### Create a Dockerfile for the image\n",
        "\n",
        "Create a Dockerfile with TorchServe as the base image by performing the following steps:\n",
        "\n",
        " - Install dependencies such as `transformers`.\n",
        " - Add model artifacts to the `/home/model-server/` directory in the container image.\n",
        " - Add the custom handler script to the `/home/model-server/` directory in the container image.\n",
        " - Create `/home/model-server/config.properties` to define the serving configuration (health and prediction listener ports).\n",
        " - Run [Torch model archiver](https://github.com/pytorch/serve/tree/master/model-archiver#creating-a-model-archive) to create a model archive file from the files copied into the `/home/model-server/` directory in the container image. The model archive is saved in the `/home/model-server/model-store/` directory with the name same as `<model-name>.mar`.\n",
        " - Launch TorchServe HTTP server that references the configuration properties and enables serving for the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6602ec14439c"
      },
      "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": "0d3993843108"
      },
      "source": [
        "#### Create a docker repository\n",
        "\n",
        "Create your own Docker repository in Artifact Registry where you push the docker image for serving predictions.\n",
        "\n",
        "1. Run the `gcloud artifacts repositories create` command to create a new Docker repository with your specified region and description.\n",
        "\n",
        "2. Run the `gcloud artifacts repositories list` command to verify that your repository is created.\n",
        "\n",
        "Set `APP_NAME` to the name of your repository."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cd1c0a3fbcdb"
      },
      "outputs": [],
      "source": [
        "# Create the repository in Artifact registry\n",
        "! gcloud artifacts repositories create {APP_NAME} --repository-format=docker --location={REGION} --description=\"Docker repository\"\n",
        "\n",
        "# List all repositories and check your repository\n",
        "! gcloud artifacts repositories list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4601a036e327"
      },
      "source": [
        "#### Build the docker image tagged with the image path\n",
        "\n",
        "Next, you build a docker image inside the created repository using Cloud Build. Cloud Build tries to locate the repository path provided in the tag.\n",
        "\n",
        "Learn more about [building and pushing a docker image with Cloud Build](https://cloud.google.com/build/docs/build-push-docker-image)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3fb45b22c008"
      },
      "outputs": [],
      "source": [
        "!gcloud builds submit --region={REGION} --tag=$CUSTOM_PREDICTOR_IMAGE_URI ./predictor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2c331bdaf370"
      },
      "source": [
        "### Deploying the serving container to Vertex AI\n",
        "\n",
        "Next, you create a model resource on Vertex AI and deploy the model to a Vertex AI Endpoint. You must deploy a model to an endpoint for serving online predictions. The deployed model runs the custom container image to serve predictions. \n",
        "\n",
        "#### Create a Vertex AI Model resource\n",
        "\n",
        "Create a Vertex AI model resource with the created model artifacts and the container image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f5a1472f735b"
      },
      "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": "a7720d771633"
      },
      "source": [
        "#### Create a Vertex AI Endpoint\n",
        "\n",
        "Create a Vertex AI Endpoint to deploy the registered Vertex AI model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fc8fe6b64dd1"
      },
      "outputs": [],
      "source": [
        "endpoint = aiplatform.Endpoint.create(display_name=endpoint_display_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "061ab3d05205"
      },
      "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:** It takes a few minutes to deploy the resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d65aae3de6c1"
      },
      "outputs": [],
      "source": [
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    deployed_model_display_name=model_display_name,\n",
        "    machine_type=DEPLOY_MACHINE_TYPE,\n",
        "    sync=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "094f7b4d8007"
      },
      "source": [
        "## Send online prediction requests\n",
        "\n",
        "Now, invoke the endpoint where the model is deployed using the Vertex AI SDK to make predictions for some test instances.\n",
        "\n",
        "### Format 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 a [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 follows:\n",
        "\n",
        "```\n",
        "[\n",
        "    {\n",
        "        \"data\": {\n",
        "            \"b64\": \"<base64 encoded string>\"\n",
        "        }\n",
        "    }\n",
        "]\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b5f0b7f8fc9a"
      },
      "outputs": [],
      "source": [
        "test_instances = [\n",
        "    b\"I went to a meeting that went really well.\",\n",
        "    b\"I ran four miles this morning with a good time.\",\n",
        "    b\"Watching the storms we had yesterday.  The lightning was incredible!\",\n",
        "    b\"The last night I said with her 'I love you '. And she said ' Yes'.\",\n",
        "    b\"I had followed a complex recipe making roasted duck, which took me hours and I had successfully made it.\",\n",
        "    b\"I woke up this morning to birds chirping.\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d0a24c94d66c"
      },
      "source": [
        "### Send online prediction requests\n",
        "\n",
        "Format the input text string, call prediction endpoint with formatted input requests and get the response."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "515acf48503c"
      },
      "outputs": [],
      "source": [
        "# print the test instances and their responses\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": "787f896a1daa"
      },
      "source": [
        "## Cleaning up \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 Job\n",
        "- Vertex AI Model\n",
        "- Vertex AI Endpoint\n",
        "- Cloud Storage Bucket (set `delete_bucket` to **True** to delete the bucket)\n",
        "- Image Regpository (Artifact Registry)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bdefede69285"
      },
      "outputs": [],
      "source": [
        "delete_bucket = False\n",
        "\n",
        "# Delete the Custom training job\n",
        "job.delete()\n",
        "\n",
        "# Undeploy the model from the endpoint\n",
        "endpoint.undeploy_all()\n",
        "# Delete the endpoint\n",
        "endpoint.delete()\n",
        "\n",
        "# Delete the Vertex AI Model resource\n",
        "model.delete()\n",
        "\n",
        "# Delete the Cloud Storage bucket\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI\n",
        "\n",
        "# Delete artifact repository\n",
        "! gcloud artifacts repositories delete $APP_NAME --location=$REGION --quiet"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "pytorch-text-sentiment-classification-custom-train-deploy.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}