import os

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# Google Cloud Notebook requires dependencies to be installed with '--user'
USER_FLAG = ""
if IS_GOOGLE_CLOUD_NOTEBOOK:
    USER_FLAG = "--user"

!pip -q install {USER_FLAG} --upgrade transformers
!pip -q install {USER_FLAG} --upgrade datasets
!pip -q install {USER_FLAG} --upgrade tqdm
!pip -q install {USER_FLAG} --upgrade cloudml-hypertune

"""
Explanation: Training, Tuning and Deploying a PyTorch Text Classification Model on Vertex AI
Fine-tuning pre-trained BERT model for sentiment classification task
Overview
This example is inspired from Token-Classification notebook and run_glue.py. 
We will be fine-tuning bert-base-cased (pre-trained) model for sentiment classification task.
You can find the details about this model at Hugging Face Hub.
For more notebooks with the state of the art PyTorch/Tensorflow/JAX, you can explore Hugging FaceNotebooks.
Dataset
We will be using IMDB movie review dataset from Hugging Face Datasets.
Objective
How to Build, Train, Tune and Deploy PyTorch models on Vertex AI and emphasize first class support for training and deploying PyTorch models on Vertex AI. 
Table of Contents
This notebook covers following sections:

Creating Notebooks instance
Training
Run Training Locally in the Notebook
Run Training Job on Vertex AI
Training with pre-built container
Training with custom container




Tuning 
Run Hyperparameter Tuning job on Vertex AI


Deploying
Deploying model on Vertex AI Predictions with custom container



Costs
This tutorial uses billable components of Google Cloud Platform (GCP):

Vertex AI Workbench
Vertex AI Training
Vertex AI Predictions
Cloud Storage
Container Registry
Cloud Build [Optional]

Learn about Vertex AI Pricing, Cloud Storage Pricing and Cloud Build Pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.
Creating Notebooks instance on Google Cloud
This notebook assumes you are working with PyTorch 1.9 DLVM development environment with GPU runtime. You can create a Notebook instance using Google Cloud Console or gcloud command.
gcloud notebooks instances create example-instance \
    --vm-image-project=deeplearning-platform-release \
    --vm-image-family=pytorch-1-9-cu110-notebooks \
    --machine-type=n1-standard-4 \
    --location=us-central1-a \
    --boot-disk-size=100 \
    --accelerator-core-count=1 \
    --accelerator-type=NVIDIA_TESLA_V100 \
    --install-gpu-driver \
    --network=default

NOTE: You must have GPU quota before you can create instances with GPUs. Check the 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. Free Trial accounts do not receive GPU quota by default.

Set up your local development environment
If you are using Colab or Google Cloud Notebooks, your environment already meets
all the requirements to run this notebook. You can skip this step.
Otherwise, make sure your environment meets this notebook's requirements.
You need the following:

The Google Cloud SDK
Git
Python 3
virtualenv
Jupyter notebook running in a virtual environment with Python 3

The Google Cloud guide to Setting up a Python development
environment and the Jupyter
installation guide provide detailed instructions
for meeting these requirements. The following steps provide a condensed set of
instructions:

Install and initialize the Cloud SDK.
Install Python 3.
Install virtualenv and create a virtual environment that uses Python 3. Activate the virtual environment.
To install Jupyter, run pip3 install jupyter on the command-line in a terminal shell.
To launch Jupyter, run jupyter notebook on the command-line in a terminal shell.
Open this notebook in the Jupyter Notebook Dashboard.

Install additional packages
Python dependencies required for this notebook are Transformers, Datasets and hypertune will be installed in the Notebooks instance itself.
End of explanation
"""


!pip -q install {USER_FLAG} --upgrade google-cloud-aiplatform

"""
Explanation: We will be using Vertex AI SDK for 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. 
Install Vertex AI SDK for Python
End of explanation
"""


# Automatically restart kernel after installs
import os

if not os.getenv("IS_TESTING"):
    # Automatically restart kernel after installs
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

"""
Explanation: Restart the Kernel
After you install the additional packages, you need to restart the notebook kernel so it can find the packages.
End of explanation
"""


PROJECT_ID = "[your-project-id]"  # <---CHANGE THIS TO YOUR PROJECT

import os

# Get your Google Cloud project ID using google.auth
if not os.getenv("IS_TESTING"):
    import google.auth

    _, PROJECT_ID = google.auth.default()
    print("Project ID: ", PROJECT_ID)

# validate PROJECT_ID
if PROJECT_ID == "" or PROJECT_ID is None or PROJECT_ID == "[your-project-id]":
    print(
        f"Please set your project id before proceeding to next step. Currently it's set as {PROJECT_ID}"
    )

"""
Explanation: Before you begin
Select a GPU runtime
Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select "Runtime --> Change runtime type > GPU"
Set up your Google Cloud project
The following steps are required, regardless of your notebook environment.

Select or create a Google Cloud project. When you first create an account, you get a $300 free credit towards your compute/storage costs.
Make sure that billing is enabled for your project.
Enable following APIs in your project required for running the tutorial
Vertex AI API
Cloud Storage API
Container Registry API
Cloud Build API


If you are running this notebook locally, you will need to install the Cloud SDK.
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.

Note: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $ into these commands.
Set your project ID
If you don't know your project ID, you may be able to get your project ID using gcloud or google.auth.
End of explanation
"""


from datetime import datetime


def get_timestamp():
    return datetime.now().strftime("%Y%m%d%H%M%S")


TIMESTAMP = get_timestamp()
print(f"TIMESTAMP = {TIMESTAMP}")

"""
Explanation: Timestamp
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.
End of explanation
"""


import os
import sys

# If you are running this notebook in Colab, run this cell and follow the
# instructions to authenticate your GCP account. This provides access to your
# Cloud Storage bucket and lets you submit training jobs and prediction
# requests.

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# If on Google Cloud Notebooks, then don't execute this code
if not IS_GOOGLE_CLOUD_NOTEBOOK:
    if "google.colab" in sys.modules:
        from google.colab import auth as google_auth

        google_auth.authenticate_user()

    # If you are running this notebook locally, replace the string below with the
    # path to your service account key and run this cell to authenticate your GCP
    # account.
    elif not os.getenv("IS_TESTING"):
        %env GOOGLE_APPLICATION_CREDENTIALS ''

"""
Explanation: Authenticate your Google Cloud account

If you are using Google Cloud Notebooks, your environment is already authenticated. Skip this step.

If you are using Colab, run the cell below and follow the instructions
when prompted to authenticate your account via oAuth.
Otherwise, follow these steps:

In the Cloud Console, go to the Create service account key page.
Click Create service account.
In the Service account name field, enter a name, and click Create.
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.
Click Create. A JSON file that contains your key downloads to your local environment.
Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell.
End of explanation
"""


BUCKET_NAME = "gs://[your-bucket-name]"  # <---CHANGE THIS TO YOUR BUCKET
REGION = "us-central1"  # @param {type:"string"}

if BUCKET_NAME == "" or BUCKET_NAME is None or BUCKET_NAME == "gs://[your-bucket-name]":
    BUCKET_NAME = f"gs://{PROJECT_ID}aip-{get_timestamp()}"

print(f"PROJECT_ID = {PROJECT_ID}")
print(f"BUCKET_NAME = {BUCKET_NAME}")
print(f"REGION = {REGION}")

"""
Explanation: Create a Cloud Storage bucket
The following steps are required, regardless of your notebook environment.
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.
Set the name of your Cloud Storage bucket below. It must be unique across all Cloud Storage buckets.
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. You may not use a Multi-Regional Storage bucket for training with Vertex AI.
End of explanation
"""


! gsutil mb -l $REGION $BUCKET_NAME

"""
Explanation: Only if your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.
End of explanation
"""


! gsutil ls -al $BUCKET_NAME

"""
Explanation: Finally, validate access to your Cloud Storage bucket by examining its contents:
End of explanation
"""


import base64
import json
import os
import random
import sys

import google.auth
from google.cloud import aiplatform
from google.cloud.aiplatform import gapic as aip
from google.cloud.aiplatform import hyperparameter_tuning as hpt
from google.protobuf.json_format import MessageToDict

from IPython.display import HTML, display

import datasets
import numpy as np
import pandas as pd
import torch
import transformers
from datasets import ClassLabel, Sequence, load_dataset
from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
                          EvalPrediction, Trainer, TrainingArguments,
                          default_data_collator)

print(f"Notebook runtime: {'GPU' if torch.cuda.is_available() else 'CPU'}")
print(f"PyTorch version : {torch.__version__}")
print(f"Transformers version : {datasets.__version__}")
print(f"Datasets version : {transformers.__version__}")

APP_NAME = "finetuned-bert-classifier"

os.environ["TOKENIZERS_PARALLELISM"] = "false"

"""
Explanation: Import libraries and define constants
End of explanation
"""


datasets = load_dataset("imdb")
datasets

"""
Explanation: Training
In this section, we will train a PyTorch model by fine-tuning pre-trained model from Hugging Face Transformers. We will train the model locally first and then on Vertex AI training service.
Training locally in the notebook
Loading the dataset
For this example we will use IMDB movie review dataset from Hugging Face Datasets for sentiment classification task. We use the Hugging Face Datasets library to download the data. This can be easily done with the function load_dataset.
End of explanation
"""


print(
    "Total # of rows in training dataset {} and size {:5.2f} MB".format(
        datasets["train"].shape[0], datasets["train"].size_in_bytes / (1024 * 1024)
    )
)
print(
    "Total # of rows in test dataset {} and size {:5.2f} MB".format(
        datasets["test"].shape[0], datasets["test"].size_in_bytes / (1024 * 1024)
    )
)

"""
Explanation: The datasets object itself is DatasetDict, which contains one key for the training, validation and test set.
End of explanation
"""


datasets["train"][0]

"""
Explanation: To access an actual element, you need to select a split first, then give an index:
End of explanation
"""


label_list = datasets["train"].unique("label")
label_list

"""
Explanation: Using the unique method to extract label list. This will allow us to experiment with other datasets without hard-coding labels.
End of explanation
"""


def show_random_elements(dataset, num_examples=2):
    assert num_examples <= len(
        dataset
    ), "Can't pick more elements than there are in the dataset."
    picks = []
    for _ in range(num_examples):
        pick = random.randint(0, len(dataset) - 1)
        while pick in picks:
            pick = random.randint(0, len(dataset) - 1)
        picks.append(pick)

    df = pd.DataFrame(dataset[picks])
    for column, typ in dataset.features.items():
        if isinstance(typ, ClassLabel):
            df[column] = df[column].transform(lambda i: typ.names[i])
        elif isinstance(typ, Sequence) and isinstance(typ.feature, ClassLabel):
            df[column] = df[column].transform(
                lambda x: [typ.feature.names[i] for i in x]
            )
    display(HTML(df.to_html()))

show_random_elements(datasets["train"])

"""
Explanation: 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).
End of explanation
"""


batch_size = 16
max_seq_length = 128
model_name_or_path = "bert-base-cased"

tokenizer = AutoTokenizer.from_pretrained(
    model_name_or_path,
    use_fast=True,
)
# 'use_fast' ensure that we use fast tokenizers (backed by Rust) from the 🤗 Tokenizers library.

"""
Explanation: Preprocessing the data
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 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.
To do all of this, we instantiate our tokenizer with the AutoTokenizer.from_pretrained method, which ensures:

we get a tokenizer that corresponds to the model architecture we want to use,
we download the vocabulary used when pre-training this specific checkpoint.

That vocabulary will be cached, so it's not downloaded again the next time we run the cell.
End of explanation
"""


tokenizer("Hello, this is one sentence!")

"""
Explanation: You can check type of models available with a fast tokenizer on the big table of models.
You can directly call this tokenizer on one sentence:
End of explanation
"""


example = datasets["train"][4]
print(example)

tokenizer(
    ["Hello", ",", "this", "is", "one", "sentence", "split", "into", "words", "."],
    is_split_into_words=True,
)

"""
Explanation: 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 if you're interested.
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:
End of explanation
"""


# Dataset loading repeated here to make this cell idempotent
# Since we are over-writing datasets variable
datasets = load_dataset("imdb")

# Mapping labels to ids
# NOTE: We can extract this automatically but the `Unique` method of the datasets
# is not reporting the label -1 which shows up in the pre-processing.
# Hence the additional -1 term in the dictionary
label_to_id = {1: 1, 0: 0, -1: 0}


def preprocess_function(examples):
    """
    Tokenize the input example texts
    NOTE: The same preprocessing step(s) will be applied
    at the time of inference as well.
    """
    args = (examples["text"],)
    result = tokenizer(
        *args, padding="max_length", max_length=max_seq_length, truncation=True
    )

    # Map labels to IDs (not necessary for GLUE tasks)
    if label_to_id is not None and "label" in examples:
        result["label"] = [label_to_id[example] for example in examples["label"]]

    return result


# apply preprocessing function to input examples
datasets = datasets.map(preprocess_function, batched=True, load_from_cache_file=True)

"""
Explanation: 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:
End of explanation
"""


model = AutoModelForSequenceClassification.from_pretrained(
    model_name_or_path, num_labels=len(label_list)
)

"""
Explanation: Fine-tuning the model
Now that our data is ready, we can download the pre-trained model and fine-tune it.

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. 
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.


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):
End of explanation
"""


args = TrainingArguments(
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=1,
    weight_decay=0.01,
    output_dir="/tmp/cls",
)

"""
Explanation: 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.
To instantiate a Trainer, we will need to define three more things. The most important is the 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:
End of explanation
"""


def compute_metrics(p: EvalPrediction):
    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
    preds = np.argmax(preds, axis=1)
    return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}

"""
Explanation: 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.
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.
End of explanation
"""


trainer = Trainer(
    model,
    args,
    train_dataset=datasets["train"],
    eval_dataset=datasets["test"],
    data_collator=default_data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

"""
Explanation: Now we create the Trainer object and we are almost ready to train.
End of explanation
"""


trainer.train()

saved_model_local_path = "./models"
!mkdir ./models

trainer.save_model(saved_model_local_path)

"""
Explanation: 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.
We can now fine-tune our model by just calling the train method:
End of explanation
"""


history = trainer.evaluate()

history

"""
Explanation: The evaluate method allows you to evaluate again on the evaluation dataset or on another dataset:
End of explanation
"""


model_name_or_path = "bert-base-cased"
label_text = {0: "Negative", 1: "Positive"}
saved_model_path = saved_model_local_path


def predict(input_text, saved_model_path):
    # initialize tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)

    # preprocess and encode input text
    tokenizer_args = (input_text,)
    predict_input = tokenizer(
        *tokenizer_args,
        padding="max_length",
        max_length=128,
        truncation=True,
        return_tensors="pt",
    )

    # load trained model
    loaded_model = AutoModelForSequenceClassification.from_pretrained(saved_model_path)

    # get predictions
    output = loaded_model(predict_input["input_ids"])

    # return labels
    label_id = torch.argmax(*output.to_tuple(), dim=1)

    print(f"Review text: {input_text}")
    print(f"Sentiment : {label_text[label_id.item()]}\n")

# example #1
review_text = (
    """Jaw dropping visual affects and action! One of the best I have seen to date."""
)
predict_input = predict(review_text, saved_model_path)

# example #2
review_text = """Take away the CGI and the A-list cast and you end up with film with less punch."""
predict_input = predict(review_text, saved_model_path)

"""
Explanation: 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.
Run predictions locally with sample examples
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 on Vertex AI Predictions.
End of explanation
"""


PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (
    "us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest"
)

PYTHON_PACKAGE_APPLICATION_DIR = "python_package"

source_package_file_name = f"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/trainer-0.1.tar.gz"
python_package_gcs_uri = (
    f"{BUCKET_NAME}/pytorch-on-gcp/{APP_NAME}/train/python_package/trainer-0.1.tar.gz"
)
python_module_name = "trainer.task"

"""
Explanation: Training on Vertex AI
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 for following reasons:

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.
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.
Training at scale: You can run a distributed training job with AI allowing you to train models in a cluster across multiple nodes in parallel and resulting in faster training time. 
Logging and Monitoring: The training service logs messages from the job to Cloud Logging and can be monitored while the job is running.

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:

STEP 1: Determine training code structure - Packaging as a Python source distribution or as a custom container image
STEP 2: Chose a custom training method - custom job, hyperparameter training job or training pipeline
STEP 3: Run the training job


Custom training methods
There are three types of Vertex AI resources you can create to train custom models on Vertex AI:

Custom jobs: 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. 
Hyperparameter tuning jobs: 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.
Training pipelines: Orchestrates custom training jobs or hyperparameter tuning jobs with additional steps after the training job is successfully completed.

Please refer to the documentation for further details.
In this notebook, we will cover Custom Jobs and Hyperparameter tuning jobs.
Packaging the training application
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.
There are two ways to package your application and dependencies and train on Vertex AI:

Create a Python source distribution with the training code and dependencies to use with a pre-built containers on Vertex AI
Use custom containers to package dependencies using Docker containers

This notebook shows both packaging options to run a custom training job on Vertex AI.
Recommended Training Application Structure
You can structure your training application in any way you like. However, the following 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.
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.
.
├── custom_container
│   ├── Dockerfile
│   ├── README.md
│   ├── scripts
│   │   └── train-cloud.sh
│   └── trainer -&gt; ../python_package/trainer/
├── python_package
│   ├── README.md
│   ├── scripts
│   │   └── train-cloud.sh
│   ├── setup.py
│   └── trainer
│       ├── __init__.py
│       ├── experiment.py
│       ├── metadata.py
│       ├── model.py
│       ├── task.py
│       └── utils.py
└── pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb    --&gt; This notebook

Main project directory contains your setup.py file or Dockerfile with the dependencies. 
Use a subdirectory named trainer to store your main application module and scripts to submit training jobs locally or cloud
Inside trainer directory:
task.py - Main application module 1) initializes and parse task arguments (hyper parameters), and 2) entry point to the trainer
model.py -  Includes function to create model with a sequence classification head from a pre-trained model.
experiment.py - Runs the model training and evaluation experiment, and exports the final model.
metadata.py - Defines metadata for classification task such as predefined model dataset name, target labels
utils.py - Includes utility functions such as data input functions to read data, save model to GCS bucket



Run Custom Job on Vertex AI Training with a pre-built container
Vertex AI provides Docker container images that can be run as pre-built containers for custom training. These containers include common dependencies used in training code based on the Machine Learning framework and framework version.
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 and package the training application code by adding standard Python dependencies - transformers, datasets and tqdm - in the setup.py file. 

Initialize the variables to define pre-built container image, location of training application and training module.
End of explanation
"""


%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/setup.py

from setuptools import find_packages
from setuptools import setup
import setuptools

from distutils.command.build import build as _build
import subprocess


REQUIRED_PACKAGES = [
    'transformers',
    'datasets',
    'tqdm',
    'cloudml-hypertune'
]

setup(
    name='trainer',
    version='0.1',
    install_requires=REQUIRED_PACKAGES,
    packages=find_packages(),
    include_package_data=True,
    description='Vertex AI | Training | PyTorch | Text Classification | Python Package'
)

"""
Explanation: 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 to include all subdirectories of the parent directory as dependencies.
End of explanation
"""


!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar

"""
Explanation: Run the following command to create a source distribution, dist/trainer-0.1.tar.gz:
End of explanation
"""


!gsutil cp {source_package_file_name} {python_package_gcs_uri}

"""
Explanation: Now upload the source distribution with training application to Cloud Storage bucket
End of explanation
"""


!gsutil ls -l {python_package_gcs_uri}

"""
Explanation: Validate the source distribution exists on Cloud Storage bucket
End of explanation
"""


!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python -m trainer.task

"""
Explanation: [Optional] Run custom training job locally
Before submitting the job to cloud, you can run the training job locally by calling the trainer.task module directly
End of explanation
"""


aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)

"""
Explanation: Run custom training job on Vertex AI
We use Vertex AI SDK for Python to create and submit training job to the Vertex AI training service.
Initialize the Vertex AI SDK for Python
End of explanation
"""


print(f"APP_NAME={APP_NAME}")
print(
    f"PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI={PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI}"
)
print(f"python_package_gcs_uri={python_package_gcs_uri}")
print(f"python_module_name={python_module_name}")

JOB_NAME = f"{APP_NAME}-pytorch-pkg-ar-{get_timestamp()}"
print(f"JOB_NAME={JOB_NAME}")

job = aiplatform.CustomPythonPackageTrainingJob(
    display_name=f"{JOB_NAME}",
    python_package_gcs_uri=python_package_gcs_uri,
    python_module_name=python_module_name,
    container_uri=PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI,
)

training_args = ["--num-epochs", "2", "--model-name", "finetuned-bert-classifier"]

model = job.run(
    replica_count=1,
    machine_type="n1-standard-8",
    accelerator_type="NVIDIA_TESLA_V100",
    accelerator_count=1,
    args=training_args,
    sync=False,
)

"""
Explanation: Configure and submit Custom Job to Vertex AI Training service
Configure a Custom Job with the pre-built container image for PyTorch and training code packaged as Python source distribution. 
NOTE: When using Vertex AI SDK for Python for submitting a training job, it creates a Training Pipeline which launches the Custom Job on Vertex AI Training service.
End of explanation
"""


job_response = MessageToDict(job._gca_resource._pb)
gcs_model_artifacts_uri = job_response["trainingTaskInputs"]["baseOutputDirectory"][
    "outputUriPrefix"
]
print(f"Model artifacts are available at {gcs_model_artifacts_uri}")

!gsutil ls -lr $gcs_model_artifacts_uri/

"""
Explanation: Monitoring progress of the Custom Job
You can monitor the custom job launched from Cloud Console following the link here or use gcloud CLI command gcloud beta ai custom-jobs stream-logs

Validate the model artifacts written to GCS by the training code after the job completes successfully
End of explanation
"""


!cd python_package && ./scripts/train-cloud.sh

"""
Explanation: [Optional] Submit custom job using gcloud CLI using Python source distribution
You can submit the training job to Vertex AI training service using gcloud beta ai custom-jobs create. gcloud command stages your training application on GCS bucket and submits the training job.
gcloud beta ai custom-jobs create \
    --display-name=${JOB_NAME} \
    --region ${REGION} \
    --python-package-uris=${PACKAGE_PATH} \
    --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/" \
    --args="--model-name","finetuned-bert-classifier","--job-dir",$JOB_DIR

worker-pool-spec parameter defines the worker pool configuration used by the custom job. Following are the fields within worker-pool-spec:
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
Set the local-package-path to the path to the training code
Set the python-module to the trainer.task which is the main module to start your application
Set the accelerator-type and machine-type to set the compute type to run the application



Refer documentation for further details.
The script at ./python_package/scripts/train-cloud.sh contains the gcloud commands to launch the custom job and monitor the logs.
End of explanation
"""


%%writefile ./custom_container/Dockerfile

# Use pytorch GPU base image
FROM us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-10:latest

# set working directory
WORKDIR /app

# Install required packages
RUN pip install google-cloud-storage transformers datasets tqdm cloudml-hypertune

# Copies the trainer code to the docker image.
COPY ./trainer/__init__.py /app/trainer/__init__.py
COPY ./trainer/experiment.py /app/trainer/experiment.py
COPY ./trainer/utils.py /app/trainer/utils.py
COPY ./trainer/metadata.py /app/trainer/metadata.py
COPY ./trainer/model.py /app/trainer/model.py
COPY ./trainer/task.py /app/trainer/task.py

# Set up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

"""
Explanation: In addition to Cloud Console, you can monitor job progress by streaming logs using gcloud CLI by passing the job id:
gcloud ai custom-jobs stream-logs &lt;job_id&gt; --region=$REGION
You can validate the model artifacts written to GCS by the training code by running the following command:
!gsutil ls -l $JOB_DIR/
Run Custom Job on Vertex AI Training with custom container
To create a training job with custom container, 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.

Build your container using Dockerfile with Training Code and Dependencies
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 using Dockerfile. We create a Dockerfile with a pre-built PyTorch container image provided by Vertex AI as the base image, install the dependencies - transformers, datasets , tqdm and cloudml-hypertune and copy the training application code.
End of explanation
"""


CUSTOM_TRAIN_IMAGE_URI = f"gcr.io/{PROJECT_ID}/pytorch_gpu_train_{APP_NAME}"

!cd ./custom_container/ && docker build -f Dockerfile -t $CUSTOM_TRAIN_IMAGE_URI ../python_package

"""
Explanation: Build the image and tag the Container Registry path (gcr.io) that you will push to.
End of explanation
"""


!docker run --gpus all -it --rm $CUSTOM_TRAIN_IMAGE_URI

"""
Explanation: [Optional] Run Training Job Locally with Custom Container
Run the container locally in detached mode to test. When running with machine with GPUs, you can use --gpus all command line flag.
End of explanation
"""


!docker push $CUSTOM_TRAIN_IMAGE_URI

"""
Explanation: Run custom training job on Vertex AI with Custom Container
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. 
NOTE: Container Registry is a central repository to store, manage, and secure your Docker container images.
Push the container to Container Registry
Push your container image with training application code and dependencies to your Container Registry.
End of explanation
"""


!gcloud container images describe $CUSTOM_TRAIN_IMAGE_URI

"""
Explanation: Validate the custom container image in Container Registry
End of explanation
"""


aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)

"""
Explanation: Initialize the Vertex AI SDK for Python
End of explanation
"""


JOB_NAME = f"{APP_NAME}-pytorch-cstm-cntr-{get_timestamp()}"

print(f"APP_NAME={APP_NAME}")
print(f"CUSTOM_TRAIN_IMAGE_URI={CUSTOM_TRAIN_IMAGE_URI}")
print(f"JOB_NAME={JOB_NAME}")

# configure the job with container image spec
job = aiplatform.CustomContainerTrainingJob(
    display_name=f"{JOB_NAME}", container_uri=f"{CUSTOM_TRAIN_IMAGE_URI}"
)

# define training code arguments
training_args = ["--num-epochs", "2", "--model-name", "finetuned-bert-classifier"]

# submit the custom job to Vertex AI training service
model = job.run(
    replica_count=1,
    machine_type="n1-standard-8",
    accelerator_type="NVIDIA_TESLA_V100",
    accelerator_count=1,
    args=training_args,
    sync=False,
)

"""
Explanation: Configure and submit Custom Job to Vertex AI Training service
Configure a Custom Job with the custom container image with training code and other dependencies
NOTE: When using Vertex AI SDK for Python for submitting a training job, it creates a Training Pipeline which launches the Custom Job to train on Vertex AI Training.
End of explanation
"""


!cd custom_container && ./scripts/train-cloud.sh

"""
Explanation: Monitoring progress of the Custom Job
You can monitor the custom job launched from Cloud Console following the link here or use gcloud CLI command gcloud beta ai custom-jobs stream-logs

[Optional] Submit Custom Job using gcloud CLI using custom container
You can submit the training job to Vertex AI training service using gcloud 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.
gcloud beta ai custom-jobs create \
    --display-name=${JOB_NAME} \
    --region ${REGION} \
    --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} \
    --args="--model-name","finetuned-bert-classifier","--job-dir",$JOB_DIR

worker-pool-spec parameter defines the worker pool configuration used by the custom job. Following are the fields within worker-pool-spec:
Set the container-image-uri to the custom container image pushed to Google Cloud Container Registry for training
Set the accelerator-type and machine-type to set the compute type to run the application



Refer documentation for further details.
The script at ./custom_container/scripts/train-cloud.sh contains the gcloud commands to launch the custom job and monitor the logs.
End of explanation
"""


!gcloud container images describe $CUSTOM_TRAIN_IMAGE_URI

"""
Explanation: In addition to Cloud Console, you can monitor job progress by streaming logs using gcloud CLI by passing the job id:
gcloud ai custom-jobs stream-logs &lt;job_id&gt; --region=$REGION
You can validate the model artifacts written to GCS by the training code by running the following command:
!gsutil ls -l $JOB_DIR/
Hyperparameter Tuning
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.
We submit a Hyperparameter Tuning job 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.

How hyperparameter tuning works in Vertex AI?
Following are the high level steps involved in running a Hyperparameter Tuning job on Vertex AI Training service:

You define the hyperparameters to tune the model along with the metric (or goal) to optimize
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. 
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. 
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

Refer to the Vertex AI documentation 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. 
Changes to training application code for hyperparameter tuning
There are few requirements to follow specific to hyperparameter tuning in Vertex AI:

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
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 to report metrics.

Previously, in the training application code to fine-tune the transformer model for sentiment analysis task, we instantiated Trainer with hyperparameters passed as training arguments (training_args). 
```
    # Estimator arguments
    args_parser.add_argument(
        '--learning-rate',
        help='Learning rate value for the optimizers.',
        default=2e-5,
        type=float)
    args_parser.add_argument(
        '--weight-decay',
        help="""
      The factor by which the learning rate should decay by the end of the
      training.
  decayed_learning_rate =
    learning_rate * decay_rate ^ (global_step / decay_steps)

  If set to 0 (default), then no decay will occur.
  If set to 0.5, then the learning rate should reach 0.5 of its original
      value at the end of the training.
  Note that decay_steps is set to train_steps.
  """,
    default=0.01,
    type=float)

# Enable hyperparameter
args_parser.add_argument(
    '--hp-tune',
    default="n",
    help='Enable hyperparameter tuning. Valida values are: "y" - enable, "n" - disable')

```
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.
```
    # set training arguments
    training_args = TrainingArguments(
        evaluation_strategy="epoch",
        learning_rate=args.learning_rate,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        num_train_epochs=args.num_epochs,
        weight_decay=args.weight_decay,
        output_dir=os.path.join("/tmp", args.model_name)
    )
# initialize our Trainer
trainer = Trainer(
    model,
    training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    data_collator=default_data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

```
To report metrics to Vertex AI when hyperparameter tuning is enabled, we call cloudml-hypertune Python package after the evaluation phase which is added as a callback 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.
```
add hyperparameter tuning callback to report metrics when enabled
if args.hp_tune == "y":
    trainer.add_callback(HPTuneCallback("accuracy", "eval_accuracy"))
class HPTuneCallback(TrainerCallback):
    """
    A custom callback class that reports a metric to hypertuner
    at the end of each epoch.
    """
def __init__(self, metric_tag, metric_value):
    super(HPTuneCallback, self).__init__()
    self.metric_tag = metric_tag
    self.metric_value = metric_value
    self.hpt = hypertune.HyperTune()

def on_evaluate(self, args, state, control, **kwargs):
    print(f"HP metric {self.metric_tag}={kwargs['metrics'][self.metric_value]}")
    self.hpt.report_hyperparameter_tuning_metric(
        hyperparameter_metric_tag=self.metric_tag,
        metric_value=kwargs['metrics'][self.metric_value],
        global_step=state.epoch)

```
Run Hyperparameter Tuning Job on Vertex AI
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.
Validate the custom container image in Container Registry
End of explanation
"""


aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)

"""
Explanation: Initialize the Vertex AI SDK for Python
End of explanation
"""


JOB_NAME = f"{APP_NAME}-pytorch-hptune-{get_timestamp()}"

print(f"APP_NAME={APP_NAME}")
print(f"CUSTOM_TRAIN_IMAGE_URI={CUSTOM_TRAIN_IMAGE_URI}")
print(f"JOB_NAME={JOB_NAME}")

"""
Explanation: Configure and submit Hyperparameter Tuning Job to Vertex AI Training service
Configure a Hyperparameter Tuning Job with the custom container image with training code and other dependencies.
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.
End of explanation
"""


training_args = [
    "--num-epochs",
    "2",
    "--model-name",
    "finetuned-bert-classifier",
    "--hp-tune",
    "y",
]

"""
Explanation: Define the training arguments with hp-tune argument set to y so that training application code can report metrics to Vertex AI
End of explanation
"""


# The spec of the worker pools including machine type and Docker image
worker_pool_specs = [
    {
        "machine_spec": {
            "machine_type": "n1-standard-8",
            "accelerator_type": "NVIDIA_TESLA_V100",
            "accelerator_count": 1,
        },
        "replica_count": 1,
        "container_spec": {"image_uri": CUSTOM_TRAIN_IMAGE_URI, "args": training_args},
    }
]

custom_job = aiplatform.CustomJob(
    display_name=JOB_NAME, worker_pool_specs=worker_pool_specs
)

"""
Explanation: Create a CustomJob with worker pool specs to define machine types, accelerators and customer container spec with the training application code
End of explanation
"""


# Dictionary representing parameters to optimize.
# The dictionary key is the parameter_id, which is passed into your training
# job as a command line argument,
# And the dictionary value is the parameter specification of the metric.
parameter_spec = {
    "learning-rate": hpt.DoubleParameterSpec(min=1e-6, max=0.001, scale="log"),
    "weight-decay": hpt.DiscreteParameterSpec(
        values=[0.0001, 0.001, 0.01, 0.1], scale=None
    ),
}

"""
Explanation: 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. 
Refer to the documentation on selecting the hyperparaneter to tune and how to define parameter specification.
End of explanation
"""


# Dictionary representing metrics to optimize.
# The dictionary key is the metric_id, which is reported by your training job,
# And the dictionary value is the optimization goal of the metric.
metric_spec = {"accuracy": "maximize"}

"""
Explanation: 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.
End of explanation
"""


hp_job = aiplatform.HyperparameterTuningJob(
    display_name=JOB_NAME,
    custom_job=custom_job,
    metric_spec=metric_spec,
    parameter_spec=parameter_spec,
    max_trial_count=5,
    parallel_trial_count=2,
    search_algorithm=None,
)

model = hp_job.run(sync=False)

"""
Explanation: Configure and submit a Hyperparameter Tuning Job with the Custom Job, metric spec, parameter spec and trial limits.

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.
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.
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. 

Refer to the documentation to understand the hyperparameter training job configuration.
End of explanation
"""


def get_trials_as_df(trials):
    results = []
    for trial in trials:
        row = {}
        t = MessageToDict(trial._pb)
        # print(t)
        row["Trial ID"], row["Status"], row["Start time"], row["End time"] = (
            t["id"],
            t["state"],
            t["startTime"],
            t.get("endTime", None),
        )

        for param in t["parameters"]:
            row[param["parameterId"]] = param["value"]

        if t["state"] == "SUCCEEDED":
            row["Training step"] = t["finalMeasurement"]["stepCount"]
            for metric in t["finalMeasurement"]["metrics"]:
                row[metric["metricId"]] = metric["value"]
        results.append(row)

    _df = pd.DataFrame(results)
    return _df

df_trials = get_trials_as_df(hp_job.trials)
df_trials

"""
Explanation: Monitoring progress of the Custom Job
You can monitor the hyperparameter tuning job launched from Cloud Console following the link here or use gcloud CLI command gcloud beta ai custom-jobs stream-logs

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
End of explanation
"""


# get trial id of the best run from the Trials
best_trial_id = df_trials.loc[df_trials["accuracy"].idxmax()]["Trial ID"]

# get base output directory where artifacts are saved
base_output_dir = MessageToDict(hp_job._gca_resource._pb)["trialJobSpec"][
    "baseOutputDirectory"
]["outputUriPrefix"]

# get the model artifacts of the best trial id
best_model_artifact_uri = f"{base_output_dir}/{best_trial_id}"

print(
    f"Model artifacts from the Hyperparameter Tuning Job with bbest trial id {best_trial_id} are located at {best_model_artifact_uri}"
)

"""
Explanation: Now from the results of Trials, you can pick the best performing Trial to deploy to Vertex AI Predictions
End of explanation
"""


!gsutil ls -r $best_model_artifact_uri/

"""
Explanation: You can validate the model artifacts written to GCS by the training code by running the following command:
End of explanation
"""


%%bash -s $BUCKET_NAME $APP_NAME

# ========================================================
# set job parameters
# ========================================================
# PROJECT_ID: Change to your project id
PROJECT_ID=$(gcloud config list --format 'value(core.project)')
        
# set job display name
JOB_PREFIX="finetuned-bert-classifier"
JOB_NAME=${JOB_PREFIX}-pytorch-hptune-$(date +%Y%m%d%H%M%S)
echo "Launching hyperparameter tuning job with display name as "$JOB_NAME

# BUCKET_NAME is a required parameter to run the cell.
BUCKET_NAME=$1

# APP_NAME: get application name
APP_NAME=$2

# JOB_DIR: Where to store prepared package and upload output model.
JOB_DIR=${BUCKET_NAME}/${JOB_PREFIX}/model/${JOB_NAME}

# custom container image URI
CUSTOM_TRAIN_IMAGE_URI='gcr.io/'${PROJECT_ID}'/pytorch_gpu_train_'${APP_NAME}

# ========================================================
# create hyperparameter tuning configuration file
# ========================================================
cat << EOF > ./python_package/hptuning_job.yaml

studySpec:
  metrics:
  - metricId: accuracy
    goal: MAXIMIZE
  parameters:
  - parameterId: learning-rate
    scaleType: UNIT_LOG_SCALE
    doubleValueSpec:
      minValue: 0.000001
      maxValue: 0.001
  - parameterId: weight-decay
    scaleType: SCALE_TYPE_UNSPECIFIED
    discreteValueSpec:
      values: [
          0.0001, 
          0.001, 
          0.01, 
          0.1
      ]
  measurementSelectionType: BEST_MEASUREMENT
trialJobSpec:
  workerPoolSpecs:
  - machineSpec:
      machineType: n1-standard-8
      acceleratorType: NVIDIA_TESLA_V100
      acceleratorCount: 1  
    replicaCount: 1
    containerSpec:
      imageUri: $CUSTOM_TRAIN_IMAGE_URI
      args: ["--num-epochs", "2", "--model-name", "finetuned-bert-classifier", "--hp-tune", "y"]
  baseOutputDirectory: 
    outputUriPrefix: $JOB_DIR/
EOF

# ========================================================
# submit hyperparameter tuning job
# ========================================================
gcloud beta ai hp-tuning-jobs create \
   --config ./python_package/hptuning_job.yaml \
   --display-name $JOB_NAME \
   --algorithm algorithm-unspecified \
   --max-trial-count 5 \
   --parallel-trial-count 2 \
   --region=us-central1

"""
Explanation: [Optional] Submit hyperparameter tuning job using gcloud CLI
You can submit the hyperparameter tuning job to Vertex AI training service using gcloud 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.
The following cell shows how to submit a hyperparameter tuning job on Vertex AI using gcloud CLI:
End of explanation
"""


%%writefile predictor/custom_handler.py

import os
import json
import logging

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from ts.torch_handler.base_handler import BaseHandler

logger = logging.getLogger(__name__)


class TransformersClassifierHandler(BaseHandler):
    """
    The handler takes an input string and returns the classification text 
    based on the serialized transformers checkpoint.
    """
    def __init__(self):
        super(TransformersClassifierHandler, self).__init__()
        self.initialized = False

    def initialize(self, ctx):
        """ Loads the model.pt file and initialized the model object.
        Instantiates Tokenizer for preprocessor to use
        Loads labels to name mapping file for post-processing inference response
        """
        self.manifest = ctx.manifest

        properties = ctx.system_properties
        model_dir = properties.get("model_dir")
        self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")

        # Read model serialize/pt file
        serialized_file = self.manifest["model"]["serializedFile"]
        model_pt_path = os.path.join(model_dir, serialized_file)
        if not os.path.isfile(model_pt_path):
            raise RuntimeError("Missing the model.pt or pytorch_model.bin file")
        
        # Load model
        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
        self.model.to(self.device)
        self.model.eval()
        logger.debug('Transformer model from path {0} loaded successfully'.format(model_dir))
        
        # Ensure to use the same tokenizer used during training
        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')

        # Read the mapping file, index to object name
        mapping_file_path = os.path.join(model_dir, "index_to_name.json")

        if os.path.isfile(mapping_file_path):
            with open(mapping_file_path) as f:
                self.mapping = json.load(f)
        else:
            logger.warning('Missing the index_to_name.json file. Inference output will default.')
            self.mapping = {"0": "Negative",  "1": "Positive"}

        self.initialized = True

    def preprocess(self, data):
        """ Preprocessing input request by tokenizing
            Extend with your own preprocessing steps as needed
        """
        text = data[0].get("data")
        if text is None:
            text = data[0].get("body")
        sentences = text.decode('utf-8')
        logger.info("Received text: '%s'", sentences)

        # Tokenize the texts
        tokenizer_args = ((sentences,))
        inputs = self.tokenizer(*tokenizer_args,
                                padding='max_length',
                                max_length=128,
                                truncation=True,
                                return_tensors = "pt")
        return inputs

    def inference(self, inputs):
        """ Predict the class of a text using a trained transformer model.
        """
        prediction = self.model(inputs['input_ids'].to(self.device))[0].argmax().item()

        if self.mapping:
            prediction = self.mapping[str(prediction)]

        logger.info("Model predicted: '%s'", prediction)
        return [prediction]

    def postprocess(self, inference_output):
        return inference_output


"""
Explanation: Deploying
Deploying a PyTorch model on Vertex AI Predictions requires to use a custom container that serves online predictions. You will deploy a container running PyTorch's TorchServe 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. 
Deploying model on Vertex AI Predictions with custom container
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 that describes the container image requirements to be compatible with Vertex AI Predictions.

Essentially, to deploy a PyTorch model on Vertex AI Predictions following are the steps:

Package the trained model artifacts including default or custom handlers by creating an archive file using Torch model archiver
Build a custom container compatible with Vertex AI Predictions to serve the model using Torchserve
Upload the model with custom container image to serve predictions as a Vertex AI Model resource
Create a Vertex AI Endpoint and deploy the model resource

Create a custom model handler to handle prediction requests
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. 
Custom handler script does the following:

Pre-process input text before sending it to the model for inference
Customize how the model is invoked for inference
Post-process output from the model before sending back a response

Please refer to the TorchServe documentation for defining a custom handler.
End of explanation
"""


%%writefile ./predictor/index_to_name.json

{
    "0": "Negative", 
    "1": "Positive"
}

"""
Explanation: Generate target label to name file [Optional]
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".
End of explanation
"""


GCS_MODEL_ARTIFACTS_URI = best_model_artifact_uri

"""
Explanation: Create custom container image to serve predictions
We will use Cloud Build to create the custom container image with following build steps:
Download model artifacts
Download model artifacts that were saved as part of the training (or hyperparameter tuning) job from Cloud Storage to local directory
End of explanation
"""


!gsutil ls -r $GCS_MODEL_ARTIFACTS_URI/model/

"""
Explanation: Validate model artifact files in the Cloud Storage bucket
End of explanation
"""


!gsutil -m cp -r $GCS_MODEL_ARTIFACTS_URI/model/ ./predictor/

!ls -ltrR ./predictor/model

"""
Explanation: Copy files from Cloud Storage to local directory
End of explanation
"""


%%bash -s $APP_NAME

APP_NAME=$1

cat << EOF > ./predictor/Dockerfile

FROM pytorch/torchserve:latest-cpu

# install dependencies
RUN python3 -m pip install --upgrade pip
RUN pip3 install transformers

USER model-server

# copy model artifacts, custom handler and other dependencies
COPY ./custom_handler.py /home/model-server/
COPY ./index_to_name.json /home/model-server/
COPY ./model/$APP_NAME/ /home/model-server/

# create torchserve configuration file
USER root
RUN printf "\nservice_envelope=json" >> /home/model-server/config.properties
RUN printf "\ninference_address=http://0.0.0.0:7080" >> /home/model-server/config.properties
RUN printf "\nmanagement_address=http://0.0.0.0:7081" >> /home/model-server/config.properties
USER model-server

# expose health and prediction listener ports from the image
EXPOSE 7080
EXPOSE 7081

# create model archive file packaging model artifacts and dependencies
RUN torch-model-archiver -f \
  --model-name=$APP_NAME \
  --version=1.0 \
  --serialized-file=/home/model-server/pytorch_model.bin \
  --handler=/home/model-server/custom_handler.py \
  --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" \
  --export-path=/home/model-server/model-store

# run Torchserve HTTP serve to respond to prediction requests
CMD ["torchserve", \
     "--start", \
     "--ts-config=/home/model-server/config.properties", \
     "--models", \
     "$APP_NAME=$APP_NAME.mar", \
     "--model-store", \
     "/home/model-server/model-store"]
EOF

echo "Writing ./predictor/Dockerfile"

"""
Explanation: Build the container image
Create a Dockerfile with TorchServe as base image:

RUN: Installs dependencies such as transformers
COPY: Add model artifacts to /home/model-server/ directory of the container image
COPY: Add custom handler script to /home/model-server/ directory of the container image
RUN: Create /home/model-server/config.properties to define the serving configuration (health and prediction listener ports)
RUN: Run Torch model archiver 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 &lt;model-name&gt;.mar
CMD: Launch Torchserve HTTP server referencing the configuration properties and enables serving for the model
End of explanation
"""


CUSTOM_PREDICTOR_IMAGE_URI = f"gcr.io/{PROJECT_ID}/pytorch_predict_{APP_NAME}"
print(f"CUSTOM_PREDICTOR_IMAGE_URI = {CUSTOM_PREDICTOR_IMAGE_URI}")

!docker build \
  --tag=$CUSTOM_PREDICTOR_IMAGE_URI \
  ./predictor

"""
Explanation: Build the docker image tagged with Container Registry (gcr.io) path
End of explanation
"""


!docker stop local_bert_classifier
!docker run -t -d --rm -p 7080:7080 --name=local_bert_classifier $CUSTOM_PREDICTOR_IMAGE_URI
!sleep 20

"""
Explanation: Run the container locally [Optional]
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

To run the container image as a container locally, run the following command:
End of explanation
"""


!curl http://localhost:7080/ping

"""
Explanation: To send the container's server a health check, run the following command:
End of explanation
"""


%%bash -s $APP_NAME

APP_NAME=$1

cat > ./predictor/instances.json <<END
{ 
   "instances": [
     { 
       "data": {
         "b64": "$(echo 'Take away the CGI and the A-list cast and you end up with film with less punch.' | base64 --wrap=0)"
       }
     }
   ]
}
END

curl -s -X POST \
  -H "Content-Type: application/json; charset=utf-8" \
  -d @./predictor/instances.json \
  http://localhost:7080/predictions/$APP_NAME/

"""
Explanation: If successful, the server returns the following response:
{
  "status": "Healthy"
}

To send the container's server a prediction request, run the following commands:
End of explanation
"""


!docker stop local_bert_classifier

"""
Explanation: This request uses a test sentence. If successful, the server returns prediction in below format:
{"predictions": ["Negative"]}

To stop the container, run the following command:
End of explanation
"""


!docker push $CUSTOM_PREDICTOR_IMAGE_URI

"""
Explanation: Deploying the serving container to Vertex AI Predictions
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. 
Push the serving container to Container Registry
Push your container image with inference code and dependencies to your Container Registry
End of explanation
"""


aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)

"""
Explanation: Initialize the Vertex AI SDK for Python
End of explanation
"""


VERSION = 1
model_display_name = f"{APP_NAME}-v{VERSION}"
model_description = "PyTorch based text classifier with custom container"

MODEL_NAME = APP_NAME
health_route = "/ping"
predict_route = f"/predictions/{MODEL_NAME}"
serving_container_ports = [7080]

model = aiplatform.Model.upload(
    display_name=model_display_name,
    description=model_description,
    serving_container_image_uri=CUSTOM_PREDICTOR_IMAGE_URI,
    serving_container_predict_route=predict_route,
    serving_container_health_route=health_route,
    serving_container_ports=serving_container_ports,
)

model.wait()

print(model.display_name)
print(model.resource_name)

"""
Explanation: Create a Model resource with custom serving container
End of explanation
"""


endpoint_display_name = f"{APP_NAME}-endpoint"
endpoint = aiplatform.Endpoint.create(display_name=endpoint_display_name)

"""
Explanation: For more context on upload or importing a model, refer documentation
Create an Endpoint for Model with Custom Container
End of explanation
"""


traffic_percentage = 100
machine_type = "n1-standard-4"
deployed_model_display_name = model_display_name
min_replica_count = 1
max_replica_count = 3
sync = True

model.deploy(
    endpoint=endpoint,
    deployed_model_display_name=deployed_model_display_name,
    machine_type=machine_type,
    traffic_percentage=traffic_percentage,
    sync=sync,
)

"""
Explanation: Deploy the Model to Endpoint
Deploying a model associates physical resources with the model so it can serve online predictions with low latency. 
NOTE: This step takes few minutes to deploy the resources.
End of explanation
"""


endpoint_display_name = f"{APP_NAME}-endpoint"
filter = f'display_name="{endpoint_display_name}"'

for endpoint_info in aiplatform.Endpoint.list(filter=filter):
    print(
        f"Endpoint display name = {endpoint_info.display_name} resource id ={endpoint_info.resource_name} "
    )

endpoint = aiplatform.Endpoint(endpoint_info.resource_name)

endpoint.list_models()

"""
Explanation: Invoking the Endpoint with deployed Model using Vertex AI SDK to make predictions
Get the Endpoint id
End of explanation
"""


test_instances = [
    b"Jaw dropping visual affects and action! One of the best I have seen to date.",
    b"Take away the CGI and the A-list cast and you end up with film with less punch.",
]

"""
Explanation: Formatting input for online prediction
This notebook uses Torchserve's KServe based inference API which is also Vertex AI Predictions compatible format. For online prediction requests, format the prediction input instances as JSON with base64 encoding as shown here:
[
    {
        "data": {
            "b64": "&lt;base64 encoded string&gt;"
        }
    }
]
Define sample texts to test predictions
End of explanation
"""


print("=" * 100)
for instance in test_instances:
    print(f"Input text: \n\t{instance.decode('utf-8')}\n")
    b64_encoded = base64.b64encode(instance)
    test_instance = [{"data": {"b64": f"{str(b64_encoded.decode('utf-8'))}"}}]
    print(f"Formatted input: \n{json.dumps(test_instance, indent=4)}\n")
    prediction = endpoint.predict(instances=test_instance)
    print(f"Prediction response: \n\t{prediction}")
    print("=" * 100)

"""
Explanation: Sending an online prediction request
Format input text string and call prediction endpoint with formatted input request and get the response
End of explanation
"""


endpoint_display_name = f"{APP_NAME}-endpoint"

%%bash -s $REGION $endpoint_display_name

REGION=$1
endpoint_display_name=$2

# get endpoint id
echo "REGION = ${REGION}"
echo "ENDPOINT DISPLAY NAME = ${endpoint_display_name}"
endpoint_id=$(gcloud beta ai endpoints list --region ${REGION} --filter "display_name=${endpoint_display_name}" --format "value(ENDPOINT_ID)")
echo "ENDPOINT_ID = ${endpoint_id}"

# call prediction endpoint
input_text="Take away the CGI and the A-list cast and you end up with film with less punch."
echo "INPUT TEXT = ${input_text}"

prediction=$(
echo """
{ 
   "instances": [
     { 
       "data": {
         "b64": "$(echo ${input_text} | base64 --wrap=0)"
       }
     }
   ]
}
""" | gcloud beta ai endpoints predict ${endpoint_id} --region=$REGION --json-request -)

echo "PREDICTION RESPONSE = ${prediction}"

"""
Explanation: [Optional] Make prediction requests using gcloud CLI
You can also call the Vertex AI Endpoint to make predictions using gcloud beta ai endpoints predict. 
The following cell shows how to make a prediction request to Vertex AI Endpoints using gcloud CLI:
End of explanation
"""


delete_custom_job = False
delete_hp_tuning_job = False
delete_endpoint = True
delete_model = False
delete_bucket = False
delete_image = False

"""
Explanation: Cleaning up
Cleaning up training and deployment resources
To clean up all Google Cloud resources used in this notebook, you can delete the Google Cloud project you used for the tutorial.
Otherwise, you can delete the individual resources you created in this tutorial:

Training Jobs
Model
Endpoint
Cloud Storage Bucket
Container Images

Set flags for the resource type to be deleted
End of explanation
"""


# API Endpoint
API_ENDPOINT = "{}-aiplatform.googleapis.com".format(REGION)

# Vertex AI location root path for your dataset, model and endpoint resources
PARENT = f"projects/{PROJECT_ID}/locations/{REGION}"

client_options = {"api_endpoint": API_ENDPOINT}

# Initialize Vertex AI SDK
aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)

# functions to create client
def create_job_client():
    client = aip.JobServiceClient(client_options=client_options)
    return client


def create_model_client():
    client = aip.ModelServiceClient(client_options=client_options)
    return client


def create_endpoint_client():
    client = aip.EndpointServiceClient(client_options=client_options)
    return client


clients = {}
clients["job"] = create_job_client()
clients["model"] = create_model_client()
clients["endpoint"] = create_endpoint_client()

"""
Explanation: Define clients for jobs, models and endpoints
End of explanation
"""


def list_custom_jobs():
    client = clients["job"]
    jobs = []
    response = client.list_custom_jobs(parent=PARENT)
    for row in response:
        _row = MessageToDict(row._pb)
        if _row["displayName"].startswith(APP_NAME):
            jobs.append((_row["name"], _row["displayName"]))
    return jobs


def list_hp_tuning_jobs():
    client = clients["job"]
    jobs = []
    response = client.list_hyperparameter_tuning_jobs(parent=PARENT)
    for row in response:
        _row = MessageToDict(row._pb)
        if _row["displayName"].startswith(APP_NAME):
            jobs.append((_row["name"], _row["displayName"]))
    return jobs


def list_models():
    client = clients["model"]
    models = []
    response = client.list_models(parent=PARENT)
    for row in response:
        _row = MessageToDict(row._pb)
        if _row["displayName"].startswith(APP_NAME):
            models.append((_row["name"], _row["displayName"]))
    return models


def list_endpoints():
    client = clients["endpoint"]
    endpoints = []
    response = client.list_endpoints(parent=PARENT)
    for row in response:
        _row = MessageToDict(row._pb)
        if _row["displayName"].startswith(APP_NAME):
            print(_row)
            endpoints.append((_row["name"], _row["displayName"]))
    return endpoints

"""
Explanation: Define functions to list the jobs, models and endpoints starting with APP_NAME defined earlier in the notebook
End of explanation
"""


# Delete the custom training using the Vertex AI fully qualified identifier for the custom training
try:
    if delete_custom_job:
        custom_jobs = list_custom_jobs()
        for job_id, job_name in custom_jobs:
            print(f"Deleting job {job_id} [{job_name}]")
            clients["job"].delete_custom_job(name=job_id)
except Exception as e:
    print(e)

"""
Explanation: Deleting custom training jobs
End of explanation
"""


# Delete the hyperparameter tuning jobs using the Vertex AI fully qualified identifier for the hyperparameter tuning job
try:
    if delete_hp_tuning_job:
        hp_tuning_jobs = list_hp_tuning_jobs()
        for job_id, job_name in hp_tuning_jobs:
            print(f"Deleting job {job_id} [{job_name}]")
            clients["job"].delete_hyperparameter_tuning_job(name=job_id)
except Exception as e:
    print(e)

"""
Explanation: Deleting hyperparameter tuning jobs
End of explanation
"""


# Delete the endpoint using the Vertex AI fully qualified identifier for the endpoint
try:
    if delete_endpoint:
        endpoints = list_endpoints()
        for endpoint_id, endpoint_name in endpoints:
            endpoint = aiplatform.Endpoint(endpoint_id)
            # undeploy models from the endpoint
            print(f"Undeploying all deployed models from the endpoint {endpoint_name}")
            endpoint.undeploy_all(sync=True)
            # deleting endpoint
            print(f"Deleting endpoint {endpoint_id} [{endpoint_name}]")
            clients["endpoint"].delete_endpoint(name=endpoint_id)
except Exception as e:
    print(e)

"""
Explanation: Undeploy models and Delete endpoints
End of explanation
"""


# Delete the model using the Vertex AI fully qualified identifier for the model
try:
    models = list_models()
    for model_id, model_name in models:
        print(f"Deleting model {model_id} [{model_name}]")
        clients["model"].delete_model(name=model_id)
except Exception as e:
    print(e)

"""
Explanation: Deleting models
End of explanation
"""


if delete_bucket and "BUCKET_NAME" in globals():
    print(f"Deleting all contents from the bucket {BUCKET_NAME}")

    shell_output = ! gsutil du -as $BUCKET_NAME
    print(
        f"Size of the bucket {BUCKET_NAME} before deleting = {shell_output[0].split()[0]} bytes"
    )

    # uncomment below line to delete contents of the bucket
    # ! gsutil rm -r $BUCKET_NAME

    shell_output = ! gsutil du -as $BUCKET_NAME
    if float(shell_output[0].split()[0]) > 0:
        print(
            "PLEASE UNCOMMENT LINE TO DELETE BUCKET. CONTENT FROM THE BUCKET NOT DELETED"
        )

    print(
        f"Size of the bucket {BUCKET_NAME} after deleting = {shell_output[0].split()[0]} bytes"
    )

"""
Explanation: Delete contents from the staging bucket

NOTE: Everything in this Cloud Storage bucket will be DELETED. Please run it with caution.
End of explanation
"""


gcr_images = !gcloud container images list --repository=gcr.io/$PROJECT_ID --filter="name~"$APP_NAME

if delete_image:
    for image in gcr_images:
        if image != "NAME":  # skip header line
            print(f"Deleting image {image} including all tags")
            !gcloud container images delete $image --force-delete-tags --quiet

"""
Explanation: Delete images from Container Registry
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.
End of explanation
"""


