[community-content] PyTorch on Google Cloud Vertex AI - Blog related notebook and scripts (#23)

* pytorch on vertex: initial commit * pytorch on vertex: reorg dir structure with new repo changes * pytorch on vertex: add cleanup script and update README files * pytorch on vertex: remove references to bucket names * PyTorch on Vertex: Updated with linter suggested changes * PyTorch on Vertex: dry-run and set resource names consistent with app name * PyTorch on Vertex: update CODEOWNERS * Pytorch on Vertex: Reorganized directory structure
2026-09-26 14:42:04 +00:00 · 2021-08-16 09:44:37 -07:00
parent bdcbaf7d1f
commit f75352aaae
28 changed files with 4858 additions and 0 deletions
@@ -0,0 +1,20 @@
 # PyTorch on Google Cloud: Text Classification
 In the PyTorch on Google Cloud series of blog posts, we aim to share how to build, train and deploy PyTorch models at scale and how to create reproducible machine learning pipelines on Google Cloud with [Vertex AI](https://cloud.google.com/vertex-ai).
 This tutorial on text classification shows how to train a PyTorch based text classification model by fine tuning a pre-trained Huggingface Transformers model and deploy the model on [Vertex AI](https://cloud.google.com/vertex-ai/docs/start/client-libraries#python) using Vertex SDK and [`gcloud ai`](https://cloud.google.com/sdk/gcloud/reference/beta/ai).
 ## Notebooks
 | <h4>Notebook</h4> | <h4>Description</h4>                       |
 | :-------- | :------- |
 | [pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb](./pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb) | Notebook to show training, hyper-parameter tuning and deploying a PyTorch model on Vertex AI |
 ## Folders
 | <h4>Folder Name</h4> | <h4>Description</h4>                       |
 | :-------- | :------- |
 | [`python_package`](./python_package) | Folder with scripts to train and tune the text classification model using PyTorch and Hugging Face Transformers. In the [notebook](./pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb), this folder is used for submitting a training job on Vertex AI using pre-built PyTorch containers. |
 | [`custom_container`](./custom_container) | Folder with reference to training scripts in [`python_package`](./python_package)folder including a `Dockerfile` to build a custom container. In the [notebook](./pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb), this folder is used for submitting a training job and hyper-parameter tuning job on Vertex AI using custom containers. |
 | [`predictor`](./predictor) | Folder with TorchServe prediction handler and Dockerfile to build a custom container with TorchServe. In the [notebook](./pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb), this folder is used for deploying a PyTorch model on Vertex AI using custom containers by running [TorchServe HTTP server](https://pytorch.org/serve/) |
@@ -0,0 +1,20 @@
 # Use pytorch GPU base image
 FROM gcr.io/cloud-aiplatform/training/pytorch-gpu.1-7
 # 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"]
@@ -0,0 +1,66 @@
 # PyTorch Custom Containers GPU Template
 ## Overview
 The directory provides code to fine tune a transformer model ([BERT-base](https://huggingface.co/bert-base-cased)) from Huggingface Transformers Library for sentiment analysis task.  [BERT](https://ai.googleblog.com/2018/11/open-sourcing-bert-state-of-art-pre.html) (Bidirectional Encoder Representations from Transformers) is a transformers model pre-trained on a large corpus of unlabeled text in a self-supervised fashion. In this sample, we use [IMDB sentiment classification dataset](https://huggingface.co/datasets/imdb) for the task. We show you packaging a PyTorch training model to submit it to Vertex AI using pre-built PyTorch containers and handling Python dependencies using [Vertex Training custom containers](https://cloud.google.com/vertex-ai/docs/training/create-custom-container?hl=hr). 
 ## Prerequisites
 * Setup your project by following the instructions from [documentation](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)
 * [Setup docker with Cloud Container Registry](https://cloud.google.com/container-registry/docs/pushing-and-pulling)
 * Change the directory to this sample and run
 `Note:` These instructions are used for local testing. When you submit a training job, no code will be executed on your local machine.
 ## Directory Structure
 * `trainer` directory: all Python modules to train the model.
 * `scripts` directory: command-line scripts to train the model on Vertex AI.
 * `setup.py`: `setup.py` scripts specifies Python dependencies required for the training job. Vertex Training uses pip to install the package on the training instances allocated for the job.
 ### Trainer Modules
 | File Name | Purpose |
 | :-------- | :------ |
 | [metadata.py](trainer/metadata.py) | Defines: metadata for classification task such as predefined model dataset name, target labels. |
 | [utils.py](trainer/utils.py) | Includes: utility functions such as data input functions to read data, save model to GCS bucket. |
 | [model.py](trainer/model.py) | Includes: function to create model with a sequence classification head from a pretrained model. |
 | [experiment.py](trainer/experiment.py) | Runs the model training and evaluation experiment, and exports the final model. |
 | [task.py](trainer/task.py) | Includes: 1) Initialize and parse task arguments (hyper parameters), and 2) Entry point to the trainer. |
 ### Scripts
 * [train-cloud.sh](scripts/train-cloud.sh) This script builds your Docker image locally, pushes the image to Container Registry and submits a custom container training job to Vertex AI.
 Please read the [documentation](https://cloud.google.com/vertex-ai/docs/training/containers-overview?hl=hr) on Vertex Training with Custom Containers for more details.
 ## How to run
 Once the prerequisites are satisfied, you may:
 1. For local testing, run (refer [notebook](../pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb) for instructions): 
    ```
    CUSTOM_TRAIN_IMAGE_URI='gcr.io/{PROJECT_ID}/pytorch_gpu_train_{APP_NAME}'
    cd ./custom_container/ && docker build -f Dockerfile -t $CUSTOM_TRAIN_IMAGE_URI ../python_package
    docker run --gpus all -it --rm $CUSTOM_TRAIN_IMAGE_URI
    ```
 2. For cloud testing, run:
    ```
    source ./scripts/train-cloud.sh
    ```
 ## Run on GPU
 The provided trainer code runs on a GPU if one is available including data loading and model creation.
 To run the trainer code on a different GPU configuration or latest PyTorch pre-built container image, make the following changes to the trainer script.
 * Update the PyTorch image URI to one of [PyTorch pre-built containers](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#available_container_images)
 * Update the [`worker-pool-spec`](https://cloud.google.com/vertex-ai/docs/training/configure-compute?hl=hr) in the gcloud command that includes a GPU
 Then, run the script to submit a Custom Job on Vertex Training job:
 ```
 source ./scripts/train-cloud.sh
 ```
 ### Versions
 This script uses the pre-built PyTorch containers for PyTorch 1.7.
 * `us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest`
@@ -0,0 +1,72 @@
 #!/bin/bash
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
 # This script performs cloud training for a PyTorch model.
 echo "Submitting PyTorch model training job to Vertex AI"
 # PROJECT_ID: Change to your project id
 PROJECT_ID=$(gcloud config list --format 'value(core.project)')
 # BUCKET_NAME: Change to your bucket name.
 BUCKET_NAME="[your-bucket-name]" # <-- CHANGE TO YOUR BUCKET NAME
 BUCKET_NAME=cloud-ai-platform-2f444b6a-a742-444b-b91a-c7519f51bd77
 # JOB_NAME: the name of your job running on AI Platform.
 JOB_PREFIX="finetuned-bert-classifier-pytorch-cstm-cntr-"
 JOB_NAME=${JOB_PREFIX}-$(date +%Y%m%d%H%M%S)-custom-job
 # This can be a GCS location to a zipped and uploaded package
 PACKAGE_PATH=./trainer
 # REGION: select a region from https://cloud.google.com/vertex-ai/docs/general/locations#available_regions
 # or use the default '`us-central1`'. The region is where the job will be run.
 REGION="us-central1"
 # JOB_DIR: Where to store prepared package and upload output model.
 JOB_DIR=gs://${BUCKET_NAME}/${JOB_PREFIX}/models/${JOB_NAME}
 # IMAGE_REPO_NAME: set a local repo name to distinquish our image
 IMAGE_REPO_NAME=pytorch_gpu_train_finetuned-bert-classifier
 # IMAGE_TAG: an easily identifiable tag for your docker image
 IMAGE_TAG=latest
 # IMAGE_URI: the complete URI location for Cloud Container Registry
 CUSTOM_TRAIN_IMAGE_URI=gcr.io/${PROJECT_ID}/${IMAGE_REPO_NAME}:${IMAGE_TAG}
 # Build the docker image
 docker build --no-cache -f Dockerfile -t $CUSTOM_TRAIN_IMAGE_URI ../python_package
 # Deploy the docker image to Cloud Container Registry
 docker push ${CUSTOM_TRAIN_IMAGE_URI}
 # Submit Custom Job to Vertex AI
 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
 echo "After the job is completed successfully, model files will be saved at $JOB_DIR/"
 # uncomment following lines to monitor the job progress by streaming logs
 # Stream the logs from the job
 # gcloud ai custom-jobs stream-logs $(gcloud ai custom-jobs list --region=$REGION --filter="displayName:"$JOB_NAME --format="get(name)")
 # # Verify the model was exported
 # echo "Verify the model was exported:"
 # gsutil ls ${JOB_DIR}/
@@ -0,0 +1 @@
 ../python_package/trainer
@@ -0,0 +1,27 @@
 FROM pytorch/torchserve:latest-cpu
 # install dependencies
 RUN pip3 install transformers
 # copy model artifacts, custom handler and other dependencies
 COPY ./custom_text_handler.py /home/model-server/
 COPY ./index_to_name.json /home/model-server/
 COPY ./model/finetuned-bert-classifier/ /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=finetuned-bert-classifier   --version=1.0   --serialized-file=/home/model-server/pytorch_model.bin   --handler=/home/model-server/custom_text_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",      "finetuned-bert-classifier=finetuned-bert-classifier.mar",      "--model-store",      "/home/model-server/model-store"]
@@ -0,0 +1,91 @@
 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 not include class name.')
        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
@@ -0,0 +1,5 @@
 {
    "0": "Negative", 
    "1": "Positive"
 }
@@ -0,0 +1,58 @@
 # PyTorch - Python Package Training
 ## Overview
 The directory provides code to fine tune a transformer model ([BERT-base](https://huggingface.co/bert-base-cased)) from Huggingface Transformers Library for sentiment analysis task.  [BERT](https://ai.googleblog.com/2018/11/open-sourcing-bert-state-of-art-pre.html) (Bidirectional Encoder Representations from Transformers) is a transformers model pre-trained on a large corpus of unlabeled text in a self-supervised fashion. In this sample, we use [IMDB sentiment classification dataset](https://huggingface.co/datasets/imdb) for the task. We show you packaging a PyTorch training model to submit it to Vertex AI using pre-built PyTorch containers and handling Python dependencies through Python build scripts (`setup.py`). 
 ## Prerequisites
 * Setup your project by following the instructions from [documentation](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)
 * Change directories to this sample.
 ## Directory Structure
 * `trainer` directory: all Python modules to train the model.
 * `scripts` directory: command-line scripts to train the model on Vertex AI.
 * `setup.py`: `setup.py` scripts specifies Python dependencies required for the training job. Vertex Training uses pip to install the package on the training instances allocated for the job.
 ### Trainer Modules
 | File Name | Purpose |
 | :-------- | :------ |
 | [metadata.py](trainer/metadata.py) | Defines: metadata for classification task such as predefined model dataset name, target labels. |
 | [utils.py](trainer/utils.py) | Includes: utility functions such as data input functions to read data, save model to GCS bucket. |
 | [model.py](trainer/model.py) | Includes: function to create model with a sequence classification head from a pretrained model. |
 | [experiment.py](trainer/experiment.py) | Runs the model training and evaluation experiment, and exports the final model. |
 | [task.py](trainer/task.py) | Includes: 1) Initialize and parse task arguments (hyper parameters), and 2) Entry point to the trainer. |
 ### Scripts
 * [train-cloud.sh](scripts/train-cloud.sh) This script submits a training job to Vertex AI
 ## How to run
 For local testing, run:
 ```
 !cd python_package && python -m trainer.task
 ```
 For cloud training, once the prerequisites are satisfied, update the
 `BUCKET_NAME` environment variable in `scripts/train-cloud.sh`. You may then
 run the following script to submit an AI Platform Training job:
 ```
 source ./python_package/scripts/train-cloud.sh
 ```
 ## Run on GPU
 The provided trainer code runs on a GPU if one is available including data loading and model creation.
 To run the trainer code on a different GPU configuration or latest PyTorch pre-built container image, make the following changes to the trainer script.
 * Update the PyTorch image URI to one of [PyTorch pre-built containers](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#available_container_images)
 * Update the [`worker-pool-spec`](https://cloud.google.com/vertex-ai/docs/training/configure-compute?hl=hr) in the gcloud command that includes a GPU
 Then, run the script to submit a Custom Job on Vertex Training job:
 ```
 source ./scripts/train-cloud.sh
 ```
 ### Versions
 This script uses the pre-built PyTorch containers for PyTorch 1.7.
 * `us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest`
@@ -0,0 +1,62 @@
 #!/bin/bash
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
 # This script performs cloud training for a PyTorch model.
 echo "Submitting Custom Job to Vertex AI to train PyTorch model"
 # BUCKET_NAME: Change to your bucket name
 BUCKET_NAME="[your-bucket-name]" # <-- CHANGE TO YOUR BUCKET NAME
 BUCKET_NAME="cloud-ai-platform-2f444b6a-a742-444b-b91a-c7519f51bd77"
 # The PyTorch image provided by Vertex AI Training.
 IMAGE_URI="us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest"
 # JOB_NAME: the name of your job running on Vertex AI.
 JOB_PREFIX="finetuned-bert-classifier-pytorch-pkg-ar-"
 JOB_NAME=${JOB_PREFIX}-$(date +%Y%m%d%H%M%S)-custom-job
 # REGION: select a region from https://cloud.google.com/vertex-ai/docs/general/locations#available_regions
 # or use the default '`us-central1`'. The region is where the job will be run.
 REGION="us-central1"
 # JOB_DIR: Where to store prepared package and upload output model.
 JOB_DIR=gs://${BUCKET_NAME}/${JOB_PREFIX}/model/${JOB_NAME}
 # validate bucket name
 if [ "${BUCKET_NAME}" = "[your-bucket-name]" ]
 then
  echo "[ERROR] INVALID VALUE: Please update the variable BUCKET_NAME with valid Cloud Storage bucket name. Exiting the script..."
  exit 1
 fi
 # Submit Custom Job to Vertex AI
 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
 echo "After the job is completed successfully, model files will be saved at $JOB_DIR/"
 # uncomment following lines to monitor the job progress by streaming logs
 # Stream the logs from the job
 # gcloud ai custom-jobs stream-logs $(gcloud ai custom-jobs list --region=$REGION --filter="displayName:"$JOB_NAME --format="get(name)")
 # # Verify the model was exported
 # echo "Verify the model was exported:"
 # gsutil ls ${JOB_DIR}/
@@ -0,0 +1,24 @@
 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'
 )
@@ -0,0 +1,134 @@
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the \"License\");
 # you may not use this file except in compliance with the License.\n",
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an \"AS IS\" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
 import numpy as np
 import hypertune
 from transformers import (
    AutoTokenizer,
    EvalPrediction,
    Trainer,
    TrainingArguments,
    default_data_collator,
    TrainerCallback
 )
 from trainer import model, metadata, utils
 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)
 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()}
 def train(args, model, train_dataset, test_dataset):
    """Create the training loop to load pretrained model and tokenizer and 
    start the training process
    Args:
      args: read arguments from the runner to set training hyperparameters
      model: The neural network that you are training
      train_dataset: The training dataset
      test_dataset: The test dataset for evaluation
    """
    # initialize the tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        metadata.PRETRAINED_MODEL_NAME,
        use_fast=True,
    )
    # 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
    )
    # add hyperparameter tuning callback to report metrics when enabled
    if args.hp_tune == "y":
        trainer.add_callback(HPTuneCallback("accuracy", "eval_accuracy"))
    # training
    trainer.train()
    return trainer
 def run(args):
    """Load the data, train, evaluate, and export the model for serving and
     evaluating.
    Args:
      args: experiment parameters.
    """
    # Open our dataset
    train_dataset, test_dataset = utils.load_data(args)
    label_list = train_dataset.unique("label")
    num_labels = len(label_list)
    # Create the model, loss function, and optimizer
    text_classifier = model.create(num_labels=num_labels)
    # Train / Test the model
    trainer = train(args, text_classifier, train_dataset, test_dataset)
    # Export the trained model
    trainer.save_model(os.path.join("/tmp", args.model_name))
    # Save the model to GCS
    if args.job_dir:
        utils.save_model(args)
    else:
        print(f"Saved model files at {os.path.join('/tmp', args.model_name)}")
        print(f"To save model files in GCS bucket, please specify job_dir starting with gs://")
@@ -0,0 +1,31 @@
 #!/usr/bin/env python
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Task type can be either 'classification', 'regression', or 'custom'.
 # This is based on the target feature in the dataset.
 TASK_TYPE = 'classification'
 # Dataset name
 DATASET_NAME = 'imdb'
 # pre-trained model name
 PRETRAINED_MODEL_NAME = 'bert-base-cased'
 # List of the class values (labels) in a classification dataset.
 TARGET_LABELS = {1:1, 0:0, -1:0}
 # maximum sequence length
 MAX_SEQ_LENGTH = 128
@@ -0,0 +1,31 @@
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the \"License\");
 # you may not use this file except in compliance with the License.\n",
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an \"AS IS\" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from transformers import AutoModelForSequenceClassification
 from trainer import metadata
 def create(num_labels):
    """create the model by loading a pretrained model or define your 
    own
    Args:
      num_labels: number of target labels
    """
    # Create the model, loss function, and optimizer
    model = AutoModelForSequenceClassification.from_pretrained(
        metadata.PRETRAINED_MODEL_NAME,
        num_labels=num_labels
    )
    return model
@@ -0,0 +1,104 @@
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the \"License\");
 # you may not use this file except in compliance with the License.\n",
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an \"AS IS\" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
 import os
 from trainer import experiment
 def get_args():
    """Define the task arguments with the default values.
    Returns:
        experiment parameters
    """
    args_parser = argparse.ArgumentParser()
    # Experiment arguments
    args_parser.add_argument(
        '--batch-size',
        help='Batch size for each training and evaluation step.',
        type=int,
        default=16)
    args_parser.add_argument(
        '--num-epochs',
        help="""\
        Maximum number of training data epochs on which to train.
        If both --train-size and --num-epochs are specified,
        --train-steps will be: (train-size/train-batch-size) * num-epochs.\
        """,
        default=1,
        type=int,
    )
    args_parser.add_argument(
        '--seed',
        help='Random seed (default: 42)',
        type=int,
        default=42,
    )
    # 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')
    # Saved model arguments
    args_parser.add_argument(
        '--job-dir',
        default=os.getenv('AIP_MODEL_DIR'),
        help='GCS location to export models')
    args_parser.add_argument(
        '--model-name',
        default="finetuned-bert-classifier",
        help='The name of your saved model')
    return args_parser.parse_args()
 def main():
    """Setup / Start the experiment
    """
    args = get_args()
    print(args)
    experiment.run(args)
 if __name__ == '__main__':
    main()
@@ -0,0 +1,99 @@
 # Copyright 2019 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the \"License\");
 # you may not use this file except in compliance with the License.\n",
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an \"AS IS\" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
 import datetime
 from google.cloud import storage
 from transformers import AutoTokenizer
 from datasets import load_dataset, load_metric, ReadInstruction
 from trainer import metadata
 def preprocess_function(examples):
    tokenizer = AutoTokenizer.from_pretrained(
        metadata.PRETRAINED_MODEL_NAME,
        use_fast=True,
    )
    # Tokenize the texts
    tokenizer_args = (
        (examples['text'],) 
    )
    result = tokenizer(*tokenizer_args, 
                       padding='max_length', 
                       max_length=metadata.MAX_SEQ_LENGTH, 
                       truncation=True)
    # TEMP: We can extract this automatically but Unique method of the dataset
    # is not reporting the label -1 which shows up in the pre-processing
    # Hence the additional -1 term in the dictionary
    label_to_id = metadata.TARGET_LABELS
    # 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[l] for l in examples["label"]]
    return result
 def load_data(args):
    """Loads the data into two different data loaders. (Train, Test)
        Args:
            args: arguments passed to the python script
    """
    # Dataset loading repeated here to make this cell idempotent
    # Since we are over-writing datasets variable
    dataset = load_dataset(metadata.DATASET_NAME)
    dataset = dataset.map(preprocess_function, 
                          batched=True, 
                          load_from_cache_file=True)
    train_dataset, test_dataset = dataset["train"], dataset["test"]
    return train_dataset, test_dataset
 def save_model(args):
    """Saves the model to Google Cloud Storage or local file system
    Args:
      args: contains name for saved model.
    """
    scheme = 'gs://'
    if args.job_dir.startswith(scheme):
        job_dir = args.job_dir.split("/")
        bucket_name = job_dir[2]
        object_prefix = "/".join(job_dir[3:]).rstrip("/")
        if object_prefix:
            model_path = '{}/{}'.format(object_prefix, args.model_name)
        else:
            model_path = '{}'.format(args.model_name)
        bucket = storage.Client().bucket(bucket_name)    
        local_path = os.path.join("/tmp", args.model_name)
        files = [f for f in os.listdir(local_path) if os.path.isfile(os.path.join(local_path, f))]
        for file in files:
            local_file = os.path.join(local_path, file)
            blob = bucket.blob("/".join([model_path, file]))
            blob.upload_from_filename(local_file)
        print(f"Saved model files in gs://{bucket_name}/{model_path}")
    else:
        print(f"Saved model files at {os.path.join('/tmp', args.model_name)}")
        print(f"To save model files in GCS bucket, please specify job_dir starting with gs://")
@@ -16,6 +16,7 @@
 # Community Content
 /community-content			@morgandu
 /community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk	@yinghsienwu
 /community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud		@RajeshThallam @ultrons