notebooks/official/sdk/SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "41aMjXA5rKoL"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_uIhhWEZdcBb"
      },
      "source": [
        "# Vertex AI SDK for Python: Custom training using Python package, managed text dataset, and TF Serving container\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/sdk/SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/sdk/SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "<a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/sdk/SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb\" target='_blank'>\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ij53c3EzrKoS"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to create a Custom Model using Custom Python Package Training, with a Vertex AI Dataset, and how to serve the model using TensorFlow-Serving Container for online prediction, and batch prediction. It requires you to provide a bucket where the dataset will be stored.\n",
        "\n",
        "Note: You may incur charges for training, prediction, storage or usage of other GCP products in connection with testing this SDK.\n",
        "\n",
        "Learn more about [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb0a1d001978"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to create a Custom Model using Custom Python Package Training and you learn how to serve the model using TensorFlow-Serving Container for online prediction. Then you perform batch prediction on the model. \n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- `Vertex AI Dataset`\n",
        "- `Veretx AI CustomPythonPackageTrainingJob`\n",
        "- `Vertex AI Model` resource\n",
        "- `Vertex AI Endpoint` resource\n",
        "- `Vertex AI Batch Prediction`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create utility functions to download data and prepare csv files for creating Vertex AI Managed    Dataset\n",
        "- Download Data\n",
        "- Prepare CSV Files for Creating Managed Dataset\n",
        "- Create Custom Training Python Package\n",
        "- Create TensorFlow Serving Container\n",
        "- Run Custom Python Package Training with Managed Text Dataset\n",
        "- Deploy a Model and Create an Endpoint on Vertex AI\n",
        "- Predict on the Endpoint\n",
        "- Create a Batch Prediction Job on the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63b0c33d2b8c"
      },
      "source": [
        "### Dataset\n",
        "#### Stack Overflow Data\n",
        "You download the stack overflow data from from  https://storage.googleapis.com/download.tensorflow.org/data/stack_overflow_16k.tar.gz and create a Vertex AI managed text dataset. \n",
        "\n",
        "The Stack Overflow Data is licensed under the Creative Commons Attribution-ShareAlike 3.0 Unported License. To view a copy of this license, visit http://creativecommons.org/licenses/by-sa/3.0/ \n",
        "\n",
        "For more information about this dataset please visit: https://console.cloud.google.com/marketplace/details/stack-exchange/stack-overflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c9a87c7a200c"
      },
      "source": [
        "### Costs \n",
        "\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "306eb9fb5ac7"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or Vertex AI Workbench Notebooks**, your environment already meets\n",
        "all the requirements to run this notebook. You can skip this step.\n",
        "\n",
        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
        "You need the following:\n",
        "\n",
        "* The Google Cloud SDK\n",
        "* Git\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps provide a condensed set of\n",
        "instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip3 install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xOMNWzTbftDr"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9eb4e79842bf"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Vertex AI Workbench Notebook product has specific requirements\n",
        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
        "    \"/opt/deeplearning/metadata/env_version\"\n",
        ")\n",
        "\n",
        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_WORKBENCH_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\"\n",
        "\n",
        "! pip3 install --upgrade google-cloud-aiplatform tensorflow {USER_FLAG} -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ed89a5527afb"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9a2bb523c478"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af017a0645ea"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
        "\n",
        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "56d591439df1"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2016f37329c0"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5bf9979b96ff"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
        "    # Get your GCP project id from gcloud\n",
        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID:\", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "09021c90b34c"
      },
      "outputs": [],
      "source": [
        "! gcloud config set project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9658ecf524b1"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. It is recommended that you choose the region closest to you.\n",
        "\n",
        "- Americas: `us-central1`\n",
        "- Europe: `europe-west4`\n",
        "- Asia Pacific: `asia-east1`\n",
        "\n",
        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
        "\n",
        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5c615e53149f"
      },
      "outputs": [],
      "source": [
        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
        "\n",
        "if REGION == \"[your-region]\":\n",
        "    REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e663bd062c6f"
      },
      "source": [
        "#### UUID\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "953fa6e5ddda"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import string\n",
        "\n",
        "\n",
        "# Generate a uuid of a specifed length(default=8)\n",
        "def generate_uuid(length: int = 8) -> str:\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
        "\n",
        "\n",
        "UUID = generate_uuid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "378e70541ba9"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
        "authenticated. \n",
        "\n",
        "**If you are using Colab**, run the cell below and follow the instructions\n",
        "when prompted to authenticate your account via oAuth.\n",
        "\n",
        "**Otherwise**, follow these steps:\n",
        "\n",
        "1. In the Cloud Console, go to the [**Create service account key**\n",
        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "\n",
        "2. Click **Create service account**.\n",
        "\n",
        "3. In the **Service account name** field, enter a name, and\n",
        "   click **Create**.\n",
        "\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\"\n",
        "into the filter box, and select\n",
        "   **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
        "\n",
        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
        "local environment.\n",
        "\n",
        "6. Enter the path to your service account key as the\n",
        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f81b13f6ed38"
      },
      "outputs": [],
      "source": [
        "# If you are running this notebook in Colab, run this cell and follow the\n",
        "# instructions to authenticate your GCP account. This provides access to your\n",
        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
        "# requests.\n",
        "\n",
        "import os\n",
        "import sys\n",
        "\n",
        "# If on Vertex AI Workbench, then don't execute this code\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
        "    \"DL_ANACONDA_HOME\"\n",
        "):\n",
        "    if \"google.colab\" in sys.modules:\n",
        "        from google.colab import auth as google_auth\n",
        "\n",
        "        google_auth.authenticate_user()\n",
        "\n",
        "    # If you are running this notebook locally, replace the string below with the\n",
        "    # path to your service account key and run this cell to authenticate your GCP\n",
        "    # account.\n",
        "    elif not os.getenv(\"IS_TESTING\"):\n",
        "        %env GOOGLE_APPLICATION_CREDENTIALS '[your-service-account-key-path]'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bba822af0b32"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "\n",
        "\n",
        "When you submit a training job using the Vertex AI SDK, you upload a Python package\n",
        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
        "the code from this package. In this tutorial, Vertex AI also saves the\n",
        "trained model that results from your job in the same bucket. Using this model artifact, you can then\n",
        "create Vertex AI model and endpoint resources in order to serve\n",
        "online predictions.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
        "Cloud Storage buckets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "52b311a64f71"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "219a24ea078b"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + UUID\n",
        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a8a62bec0259"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "91c46850b49b"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4e69d430073b"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "835eaacd691f"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2d86f10ef734"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d630a1be6a5f"
      },
      "outputs": [],
      "source": [
        "import csv\n",
        "import os\n",
        "\n",
        "from google.cloud import aiplatform, storage\n",
        "from tensorflow.keras import utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YHNj4P3vrKoV"
      },
      "source": [
        "### Set Your Application Name, Task Name, and Directories.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RnNL5St8rKoV"
      },
      "outputs": [],
      "source": [
        "APP_NAME = \"keras-text-class-stack-overflow-tag\"\n",
        "TASK_TYPE = \"mbsdk_custom-py-pkg-training\"\n",
        "\n",
        "TASK_NAME = f\"{TASK_TYPE}_{APP_NAME}\"\n",
        "\n",
        "TASK_DIR = f\"./{TASK_NAME}\"\n",
        "DATA_DIR = f\"{TASK_DIR}/data\"\n",
        "\n",
        "print(f\"Task Name:      {TASK_NAME}\")\n",
        "print(f\"Task Directory: {TASK_DIR}\")\n",
        "print(f\"Data Directory: {DATA_DIR}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9hZmDcfLS13W"
      },
      "source": [
        "### Set a GCS prefix\n",
        "\n",
        "If you want to centeralize all input and output files under the gcs location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W_4gzaP0SyQd"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
        "GCS_PREFIX = f\"{TASK_TYPE}/{APP_NAME}\"\n",
        "\n",
        "print(f\"Bucket Name:    {BUCKET_NAME}\")\n",
        "print(f\"GCS Prefix:     {GCS_PREFIX}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DJF047yNftDw"
      },
      "source": [
        "### Utility functions to download data and prepare csv files for creating Vertex AI managed dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9yOl-l_oftDx"
      },
      "outputs": [],
      "source": [
        "def upload_blob(bucket_name, source_file_name, destination_blob_name):\n",
        "    \"\"\"Uploads a file to the bucket.\"\"\"\n",
        "\n",
        "    storage_client = storage.Client()\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "    blob = bucket.blob(destination_blob_name)\n",
        "\n",
        "    blob.upload_from_filename(source_file_name)\n",
        "\n",
        "    destination_file_name = os.path.join(\"gs://\", bucket_name, destination_blob_name)\n",
        "\n",
        "    return destination_file_name\n",
        "\n",
        "\n",
        "def download_data(data_dir):\n",
        "    \"\"\"Download data.\"\"\"\n",
        "\n",
        "    if not os.path.exists(data_dir):\n",
        "        os.makedirs(data_dir)\n",
        "\n",
        "    url = \"https://storage.googleapis.com/download.tensorflow.org/data/stack_overflow_16k.tar.gz\"\n",
        "    dataset = utils.get_file(\n",
        "        \"stack_overflow_16k.tar.gz\",\n",
        "        url,\n",
        "        untar=True,\n",
        "        cache_dir=data_dir,\n",
        "        cache_subdir=\"\",\n",
        "    )\n",
        "    data_dir = os.path.join(os.path.dirname(dataset))\n",
        "\n",
        "    return data_dir\n",
        "\n",
        "\n",
        "def upload_train_data_to_gcs(train_data_dir, bucket_name, destination_blob_prefix):\n",
        "    \"\"\"Create CSV file using train data content.\"\"\"\n",
        "\n",
        "    train_data_dir = os.path.join(data_dir, \"train\")\n",
        "    train_data_fn = os.path.join(data_dir, \"train.csv\")\n",
        "\n",
        "    fp = open(train_data_fn, \"w\", encoding=\"utf8\")\n",
        "    writer = csv.writer(\n",
        "        fp, delimiter=\",\", quotechar='\"', quoting=csv.QUOTE_ALL, lineterminator=\"\\n\"\n",
        "    )\n",
        "\n",
        "    for root, _, files in os.walk(train_data_dir):\n",
        "        for file in files:\n",
        "            if file.endswith(\".txt\"):\n",
        "                class_name = root.split(\"/\")[-1]\n",
        "                file_fn = os.path.join(root, file)\n",
        "                with open(file_fn) as f:\n",
        "                    content = f.readlines()\n",
        "                    lines = [x.strip().strip('\"') for x in content]\n",
        "                    writer.writerow((lines[0], class_name))\n",
        "\n",
        "    fp.close()\n",
        "\n",
        "    train_gcs_url = upload_blob(\n",
        "        bucket_name, train_data_fn, os.path.join(destination_blob_prefix, \"train.csv\")\n",
        "    )\n",
        "\n",
        "    return train_gcs_url"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q3UDGt8MrKoY"
      },
      "source": [
        "### Download data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ks2ahWfUrKoa"
      },
      "outputs": [],
      "source": [
        "data_dir = download_data(DATA_DIR)\n",
        "print(f\"Data is downloaded to: {data_dir}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l7AHtjj3rKoa"
      },
      "outputs": [],
      "source": [
        "!ls $data_dir"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "puroD_Wr-W4Y"
      },
      "outputs": [],
      "source": [
        "!ls $data_dir/train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQNL5AjtrKoa"
      },
      "source": [
        "### Prepare csv files for creating managed dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ewrG6VArKob"
      },
      "source": [
        "#### Create csv files using data content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GvbpGr0ArKob"
      },
      "outputs": [],
      "source": [
        "gcs_source_train_url = upload_train_data_to_gcs(\n",
        "    train_data_dir=os.path.join(data_dir, \"train\"),\n",
        "    bucket_name=BUCKET_NAME,\n",
        "    destination_blob_prefix=f\"{GCS_PREFIX}/data\",\n",
        ")\n",
        "\n",
        "print(f\"Train data content is loaded to {gcs_source_train_url}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gg_OiKJUrKoc"
      },
      "outputs": [],
      "source": [
        "!gsutil ls gs://$BUCKET_NAME/$GCS_PREFIX/data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qzKx7Z2WrKoc"
      },
      "source": [
        "# Create Custom Training Python Package\n",
        "\n",
        "Before you can perform custom training with a pre-built container, you must create a [Python Source Distribution](https://docs.python.org/3/distutils/sourcedist.html) that contains your training application and upload it to a Cloud Storage bucket that your Google Cloud project can access.\n",
        "\n",
        "You create a directory and write all of our package build artifacts into that folder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eBDjNpkLrKoc"
      },
      "outputs": [],
      "source": [
        "PYTHON_PACKAGE_APPLICATION_DIR = f\"{TASK_NAME}/trainer\"\n",
        "\n",
        "!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR\n",
        "!touch $PYTHON_PACKAGE_APPLICATION_DIR/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sPSRvSvMrKod"
      },
      "source": [
        "### Write the Training Script"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zupq7mwzrKod"
      },
      "outputs": [],
      "source": [
        "%%writefile {PYTHON_PACKAGE_APPLICATION_DIR}/task.py\n",
        "\n",
        "\n",
        "import os\n",
        "import argparse\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras import losses\n",
        "from tensorflow.keras.layers.experimental.preprocessing import TextVectorization\n",
        "\n",
        "import json\n",
        "import tqdm\n",
        "\n",
        "VOCAB_SIZE = 10000\n",
        "MAX_SEQUENCE_LENGTH = 250\n",
        "\n",
        "def str2bool(v):\n",
        "  if isinstance(v, bool):\n",
        "    return v\n",
        "  if v.lower() in ('yes', 'true', 't', 'y', '1'):\n",
        "    return True\n",
        "  elif v.lower() in ('no', 'false', 'f', 'n', '0'):\n",
        "    return False\n",
        "  else:\n",
        "    raise argparse.ArgumentTypeError('Boolean value expected.')\n",
        "\n",
        "def build_model(num_classes, loss, optimizer, metrics, vectorize_layer):\n",
        "  # vocab_size is VOCAB_SIZE + 1 since 0 is used additionally for padding.\n",
        "  model = tf.keras.Sequential([\n",
        "      vectorize_layer,\n",
        "      layers.Embedding(VOCAB_SIZE + 1, 64, mask_zero=True),\n",
        "      layers.Conv1D(64, 5, padding=\"valid\", activation=\"relu\", strides=2),\n",
        "      layers.GlobalMaxPooling1D(),\n",
        "      layers.Dense(num_classes),\n",
        "      layers.Activation('softmax')\n",
        "  ])\n",
        "  model.compile(\n",
        "      loss=loss,\n",
        "      optimizer=optimizer,\n",
        "      metrics=metrics)\n",
        "\n",
        "  return model\n",
        "\n",
        "def get_string_labels(predicted_scores_batch, class_names):\n",
        "  predicted_labels = tf.argmax(predicted_scores_batch, axis=1)\n",
        "  predicted_labels = tf.gather(class_names, predicted_labels)\n",
        "  return predicted_labels\n",
        "\n",
        "def predict(export_model, class_names, inputs):\n",
        "  predicted_scores = export_model.predict(inputs)\n",
        "  predicted_labels = get_string_labels(predicted_scores, class_names)\n",
        "  return predicted_labels\n",
        "\n",
        "def parse_args():\n",
        "  parser = argparse.ArgumentParser(\n",
        "      description='Keras Text Classification on Stack Overflow Questions')\n",
        "  parser.add_argument(\n",
        "      '--epochs', default=25, type=int, help='number of training epochs')\n",
        "  parser.add_argument(\n",
        "      '--batch-size', default=16, type=int, help='mini-batch size')\n",
        "  parser.add_argument(\n",
        "      '--model-dir', default=os.getenv('AIP_MODEL_DIR'), type=str, help='model directory')\n",
        "  parser.add_argument(\n",
        "      '--data-dir', default='./data', type=str, help='data directory')\n",
        "  parser.add_argument(\n",
        "      '--test-run', default=False, type=str2bool, help='test run the training application, i.e. 1 epoch for training using sample dataset')\n",
        "  parser.add_argument(\n",
        "      '--model-version', default=1, type=int, help='model version')\n",
        "  args = parser.parse_args()\n",
        "  return args\n",
        "\n",
        "def load_aip_dataset(aip_data_uri_pattern, batch_size, class_names, test_run, shuffle=True, seed=42):\n",
        "\n",
        "  data_file_urls = list()\n",
        "  labels = list()\n",
        "\n",
        "  class_indices = dict(zip(class_names, range(len(class_names))))\n",
        "  num_classes = len(class_names)\n",
        "\n",
        "  for aip_data_uri in tqdm.tqdm(tf.io.gfile.glob(pattern=aip_data_uri_pattern)):\n",
        "    with tf.io.gfile.GFile(name=aip_data_uri, mode='r') as gfile:\n",
        "      for line in gfile.readlines():\n",
        "        line = json.loads(line)\n",
        "        data_file_urls.append(line['textContent'])\n",
        "        classification_annotation = line['classificationAnnotations'][0]\n",
        "        label = classification_annotation['displayName']\n",
        "        labels.append(class_indices[label])\n",
        "        if test_run:\n",
        "          break\n",
        "\n",
        "  data = list()\n",
        "  for data_file_url in tqdm.tqdm(data_file_urls):\n",
        "    with tf.io.gfile.GFile(name=data_file_url, mode='r') as gf:\n",
        "      txt = gf.read()\n",
        "      data.append(txt)\n",
        "\n",
        "  print(f' data files count: {len(data_file_urls)}')\n",
        "  print(f' data count: {len(data)}')\n",
        "  print(f' labels count: {len(labels)}')\n",
        "\n",
        "  dataset = tf.data.Dataset.from_tensor_slices(data)\n",
        "  label_ds = tf.data.Dataset.from_tensor_slices(labels)\n",
        "  label_ds = label_ds.map(lambda x: tf.one_hot(x, num_classes))\n",
        "\n",
        "  dataset = tf.data.Dataset.zip((dataset, label_ds))\n",
        "\n",
        "  if shuffle:\n",
        "    # Shuffle locally at each iteration\n",
        "    dataset = dataset.shuffle(buffer_size=batch_size * 8, seed=seed)\n",
        "  dataset = dataset.batch(batch_size)\n",
        "  # Users may need to reference `class_names`.\n",
        "  dataset.class_names = class_names\n",
        "\n",
        "  return dataset\n",
        "\n",
        "def main():\n",
        "\n",
        "  args = parse_args()\n",
        "\n",
        "  class_names = ['csharp', 'java', 'javascript', 'python']\n",
        "  class_indices = dict(zip(class_names, range(len(class_names))))\n",
        "  num_classes = len(class_names)\n",
        "  print(f' class names: {class_names}')\n",
        "  print(f' class indices: {class_indices}')\n",
        "  print(f' num classes: {num_classes}')\n",
        "\n",
        "  epochs = 1 if args.test_run else args.epochs\n",
        "\n",
        "  aip_model_dir = os.environ.get('AIP_MODEL_DIR')\n",
        "  aip_data_format = os.environ.get('AIP_DATA_FORMAT')\n",
        "  aip_training_data_uri = os.environ.get('AIP_TRAINING_DATA_URI')\n",
        "  aip_validation_data_uri = os.environ.get('AIP_VALIDATION_DATA_URI')\n",
        "  aip_test_data_uri = os.environ.get('AIP_TEST_DATA_URI')\n",
        "\n",
        "  print(f\"aip_model_dir: {aip_model_dir}\")\n",
        "  print(f\"aip_data_format: {aip_data_format}\")\n",
        "  print(f\"aip_training_data_uri: {aip_training_data_uri}\")\n",
        "  print(f\"aip_validation_data_uri: {aip_validation_data_uri}\")\n",
        "  print(f\"aip_test_data_uri: {aip_test_data_uri}\")\n",
        "\n",
        "  print('Loading AIP dataset')\n",
        "  train_ds = load_aip_dataset(\n",
        "      aip_training_data_uri, args.batch_size, class_names, args.test_run)\n",
        "  print('AIP training dataset is loaded')\n",
        "  val_ds = load_aip_dataset(\n",
        "      aip_validation_data_uri, 1, class_names, args.test_run)\n",
        "  print('AIP validation dataset is loaded')\n",
        "  test_ds = load_aip_dataset(\n",
        "      aip_test_data_uri, 1, class_names, args.test_run)\n",
        "  print('AIP test dataset is loaded')\n",
        "\n",
        "  vectorize_layer = TextVectorization(\n",
        "      max_tokens=VOCAB_SIZE,\n",
        "      output_mode='int',\n",
        "      output_sequence_length=MAX_SEQUENCE_LENGTH)\n",
        "\n",
        "  train_text = train_ds.map(lambda text, labels: text)\n",
        "  vectorize_layer.adapt(train_text)\n",
        "  print('The vectorize_layer is adapted')\n",
        "\n",
        "\n",
        "  print('Build model')\n",
        "  optimizer = 'adam'\n",
        "  metrics = ['accuracy']\n",
        "\n",
        "  model = build_model(\n",
        "      num_classes, losses.CategoricalCrossentropy(from_logits=True), optimizer, metrics, vectorize_layer)\n",
        "\n",
        "  history = model.fit(train_ds, validation_data=val_ds, epochs=epochs)\n",
        "  history = history.history\n",
        "\n",
        "  print('Training accuracy: {acc}, loss: {loss}'.format(\n",
        "      acc=history['accuracy'][-1], loss=history['loss'][-1]))\n",
        "  print('Validation accuracy: {acc}, loss: {loss}'.format(\n",
        "      acc=history['val_accuracy'][-1], loss=history['val_loss'][-1]))\n",
        "\n",
        "  loss, accuracy = model.evaluate(test_ds)\n",
        "  print('Test accuracy: {acc}, loss: {loss}'.format(\n",
        "      acc=accuracy, loss=loss))\n",
        "\n",
        "  inputs = [\n",
        "      \"how do I extract keys from a dict into a list?\",  # python\n",
        "      \"debug public static void main(string[] args) {...}\",  # java\n",
        "  ]\n",
        "  predicted_labels = predict(model, class_names, inputs)\n",
        "  for input, label in zip(inputs, predicted_labels):\n",
        "    print(f'Question: {input}')\n",
        "    print(f'Predicted label: {label.numpy()}')\n",
        "\n",
        "  model_export_path = os.path.join(args.model_dir, str(args.model_version))\n",
        "  model.save(model_export_path)\n",
        "  print(f'Model version {args.model_version} is exported to {args.model_dir}')\n",
        "\n",
        "  loaded = tf.saved_model.load(model_export_path)\n",
        "  input_name = list(loaded.signatures['serving_default'].structured_input_signature[1].keys())[0]\n",
        "  print(f'Serving function input: {input_name}')\n",
        "\n",
        "  return\n",
        "\n",
        "if __name__ == '__main__':\n",
        "  main()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAhynoNBrKoh"
      },
      "source": [
        "### Build Package"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kXSBQcqdrKof"
      },
      "outputs": [],
      "source": [
        "%%writefile {TASK_DIR}/setup.py\n",
        "\n",
        "from setuptools import find_packages\n",
        "from setuptools import setup\n",
        "\n",
        "setup(\n",
        "    name='trainer',\n",
        "    version='0.1',\n",
        "    packages=find_packages(),\n",
        "    install_requires=(),\n",
        "    include_package_data=True,\n",
        "    description='My training application.'\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MclanW0UrKoh"
      },
      "outputs": [],
      "source": [
        "!ls $TASK_DIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j19Nk2vzrKoh"
      },
      "outputs": [],
      "source": [
        "!cd $TASK_DIR && python3 setup.py sdist --formats=gztar"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9c1r_k69rKoi"
      },
      "outputs": [],
      "source": [
        "!ls -ltr $TASK_DIR/dist/trainer-0.1.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zyQv9AcNrKoi"
      },
      "source": [
        "### Upload the Package to GCS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9WjF8SFrrKoi"
      },
      "outputs": [],
      "source": [
        "destination_blob_name = f\"custom-training-python-package/{APP_NAME}/trainer-0.1.tar.gz\"\n",
        "source_file_name = f\"{TASK_DIR}/dist/trainer-0.1.tar.gz\"\n",
        "\n",
        "python_package_gcs_uri = upload_blob(\n",
        "    BUCKET_NAME, source_file_name, destination_blob_name\n",
        ")\n",
        "python_module_name = \"trainer.task\"\n",
        "\n",
        "print(f\"Custom Training Python Package is uploaded to: {python_package_gcs_uri}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWJpv4ejrWi0"
      },
      "source": [
        "# Create TensorFlow Serving container"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pI50mO1bPMS1"
      },
      "outputs": [],
      "source": [
        "TF_SERVING_CONTAINER_IMAGE_URI = f\"gcr.io/{PROJECT_ID}/tf-serving\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IMOad4Z4aOgv"
      },
      "source": [
        "Download the TensorFlow Serving Docker image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KC8e1yrHrVyA"
      },
      "outputs": [],
      "source": [
        "if not IS_COLAB:\n",
        "    !docker pull tensorflow/serving:2.8.0\n",
        "else:\n",
        "    # install docker daemon\n",
        "    ! apt-get -qq install docker.io"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "33d70b04763b"
      },
      "source": [
        "Configure docker authentication with Container Registry\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f729a8ac2c1f"
      },
      "outputs": [],
      "source": [
        "! gcloud auth configure-docker gcr.io --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b7599a2703e5"
      },
      "source": [
        "Create a tag for registering the image and register the image with Cloud Container Registry (gcr.io)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bw_A7ynTsjTF"
      },
      "outputs": [],
      "source": [
        "if not IS_COLAB:\n",
        "    !docker tag tensorflow/serving:2.8.0 $TF_SERVING_CONTAINER_IMAGE_URI\n",
        "    !docker push $TF_SERVING_CONTAINER_IMAGE_URI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "86_VILTZu-Ck"
      },
      "outputs": [],
      "source": [
        "%%bash -s $IS_COLAB $TF_SERVING_CONTAINER_IMAGE_URI\n",
        "if [ $1 == \"False\" ]; then\n",
        "  exit 0\n",
        "fi\n",
        "set -x\n",
        "dockerd -b none --iptables=0 -l warn &\n",
        "for i in $(seq 5); do [ ! -S \"/var/run/docker.sock\" ] && sleep 2 || break; done\n",
        "docker pull tensorflow/serving:2.8.0\n",
        "docker tag tensorflow/serving:2.8.0 $2\n",
        "docker push $2\n",
        "kill $(jobs -p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rc_MIJwlrKoj"
      },
      "source": [
        "# Run custom Python package training with managed text dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eiNL0HSVrKoj"
      },
      "source": [
        "## Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the *client* for Vertex AI SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J5MqvGtMrKoj"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KGT1uT-HrKoj"
      },
      "source": [
        "## Create a Vertex AI dataset resource\n",
        "You create a Vertex AI text dataset using the previously prepared csv files. Choose one of the options below. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lY566K2crKok"
      },
      "outputs": [],
      "source": [
        "dataset_display_name = f\"temp-{APP_NAME}-content\"\n",
        "gcs_source = gcs_source_train_url"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kyLoUsx9rKok"
      },
      "source": [
        "####  Create a dataset with csv file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h7QwC1ThrKok"
      },
      "outputs": [],
      "source": [
        "dataset = aiplatform.TextDataset.create(\n",
        "    display_name=dataset_display_name,\n",
        "    gcs_source=gcs_source,\n",
        "    import_schema_uri=aiplatform.schema.dataset.ioformat.text.single_label_classification,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NzLvubIrrKon"
      },
      "source": [
        "## Launch a training job and create a model on Vertex AI\n",
        "\n",
        "Next, you train a model with the Python package you just built."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cyz86ixRYNtT"
      },
      "source": [
        "### Config a Training Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9zk8SRuK1Xvi"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = APP_NAME\n",
        "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rxc_yVxkbbbd"
      },
      "source": [
        "You need to specify the Python package that was built and uploaded to GCS, the module name of the Python package, the pre-built training container image uri for training, and in this example,  are using TensorFlow serving container for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uRGrFdxOftD1"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=f\"temp_{TASK_NAME}_tf-serving\",\n",
        "    python_package_gcs_uri=python_package_gcs_uri,\n",
        "    python_module_name=python_module_name,\n",
        "    container_uri=PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI,\n",
        "    model_serving_container_image_uri=TF_SERVING_CONTAINER_IMAGE_URI,\n",
        "    model_serving_container_command=[\"/usr/bin/tensorflow_model_server\"],\n",
        "    model_serving_container_args=[\n",
        "        f\"--model_name={MODEL_NAME}\",\n",
        "        \"--model_base_path=$(AIP_STORAGE_URI)\",\n",
        "        \"--rest_api_port=8080\",\n",
        "        \"--port=8500\",\n",
        "        \"--file_system_poll_wait_seconds=31540000\",\n",
        "    ],\n",
        "    model_serving_container_predict_route=f\"/v1/models/{MODEL_NAME}:predict\",\n",
        "    model_serving_container_health_route=f\"/v1/models/{MODEL_NAME}\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPfZGlZeTN72"
      },
      "source": [
        "### Run the Training Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V6iY5wcerKon"
      },
      "outputs": [],
      "source": [
        "model = job.run(\n",
        "    dataset=dataset,\n",
        "    annotation_schema_uri=aiplatform.schema.dataset.annotation.text.classification,\n",
        "    args=[\"--epochs\", \"50\"],\n",
        "    replica_count=1,\n",
        "    model_display_name=f\"temp_{TASK_NAME}_tf-serving\",\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zV1PjANLrKoo"
      },
      "outputs": [],
      "source": [
        "model.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XO0-zdXxrKoo"
      },
      "source": [
        "# Deploy a Model and Create an Endpoint on Vertex AI\n",
        "\n",
        "Deploy your model, then wait until the model FINISHES deployment before proceeding to prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tEg2IDwPftD2"
      },
      "outputs": [],
      "source": [
        "endpoint = model.deploy(machine_type=\"n1-standard-4\", sync=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9AoIJdCHrKop"
      },
      "outputs": [],
      "source": [
        "endpoint.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z87MXJEDrKoq"
      },
      "source": [
        "## Predict on the Endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vQXGCg8IrKoq"
      },
      "outputs": [],
      "source": [
        "class_names = [\"csharp\", \"java\", \"javascript\", \"python\"]\n",
        "\n",
        "class_ids = range(len(class_names))\n",
        "\n",
        "class_indices = dict(zip(class_names, class_ids))\n",
        "class_maps = dict(zip(class_ids, class_names))\n",
        "print(f\"Class Indices: {class_indices}\")\n",
        "print(f\"Class Maps:    {class_maps}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Di8RtRxipm63"
      },
      "outputs": [],
      "source": [
        "text_inputs = [\n",
        "    \"how do I extract keys from a dict into a list?\",  # python\n",
        "    \"debug public static void main(string[] args) {...}\",  # java\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dPHYGAlvlTAj"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "predictions = endpoint.predict(instances=text_inputs)\n",
        "for text, predicted_scores in zip(text_inputs, predictions.predictions):\n",
        "    class_id = np.argmax(predicted_scores)\n",
        "    class_name = class_maps[class_id]\n",
        "    print(f\"Question: {text}\")\n",
        "    print(f\"Predicted Tag: {class_name}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gRkvzr2-LgA9"
      },
      "source": [
        "# Batch Prediction Job on the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlzRioF7UGFd"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "def upload_test_data_to_gcs(test_data_dir, test_gcs_url):\n",
        "    \"\"\"Create JSON file using test data content.\"\"\"\n",
        "\n",
        "    input_name = \"text_vectorization_input\"\n",
        "\n",
        "    with tf.io.gfile.GFile(test_gcs_url, \"w\") as gf:\n",
        "\n",
        "        for root, _, files in os.walk(test_data_dir):\n",
        "            for file in files:\n",
        "                if file.endswith(\".txt\"):\n",
        "                    file_fn = os.path.join(root, file)\n",
        "                    with open(file_fn) as f:\n",
        "                        content = f.readlines()\n",
        "                        lines = [x.strip().strip('\"') for x in content]\n",
        "\n",
        "                        data = {input_name: lines[0]}\n",
        "                        gf.write(json.dumps(data))\n",
        "                        gf.write(\"\\n\")\n",
        "    return"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C_sbOm-5ud5C"
      },
      "outputs": [],
      "source": [
        "gcs_source_test_url = f\"gs://{BUCKET_NAME}/{GCS_PREFIX}/data/test.json\"\n",
        "upload_test_data_to_gcs(\n",
        "    test_data_dir=os.path.join(data_dir, \"test\"), test_gcs_url=gcs_source_test_url\n",
        ")\n",
        "\n",
        "print(f\"Test data content is loaded to {gcs_source_test_url}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QBMBk2WxLqBP"
      },
      "outputs": [],
      "source": [
        "!gsutil ls $gcs_source_test_url"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2JhoTiD5LuDW"
      },
      "outputs": [],
      "source": [
        "batch_predict_job = model.batch_predict(\n",
        "    job_display_name=f\"temp_{TASK_NAME}_tf-serving\",\n",
        "    gcs_source=gcs_source_test_url,\n",
        "    gcs_destination_prefix=f\"gs://{BUCKET_NAME}/{GCS_PREFIX}/batch_prediction\",\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "retsI3LLls_W"
      },
      "outputs": [],
      "source": [
        "batch_predict_job.wait()\n",
        "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
        "\n",
        "prediction_errors_stats = list()\n",
        "prediction_results = list()\n",
        "for blob in bp_iter_outputs:\n",
        "    if blob.name.split(\"/\")[-1].startswith(\"prediction.errors_stats\"):\n",
        "        prediction_errors_stats.append(blob.name)\n",
        "    if blob.name.split(\"/\")[-1].startswith(\"prediction.results\"):\n",
        "        prediction_results.append(blob.name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KmfK3Tzzlv9C"
      },
      "outputs": [],
      "source": [
        "tags = list()\n",
        "for prediction_result in prediction_results:\n",
        "    gfile_name = f\"gs://{bp_iter_outputs.bucket.name}/{prediction_result}\"\n",
        "    with tf.io.gfile.GFile(name=gfile_name, mode=\"r\") as gfile:\n",
        "        for line in gfile.readlines():\n",
        "            line = json.loads(line)\n",
        "            text = line[\"instance\"][\"text_vectorization_input\"][0]\n",
        "            prediction = line[\"prediction\"]\n",
        "            class_id = np.argmax(prediction)\n",
        "            class_name = class_maps[class_id]\n",
        "            tags.append([text, class_name])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UcMQ-daLn_oh"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "tags_df = pd.DataFrame(tags, columns=[\"question\", \"tag\"])\n",
        "tags_df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QUV_PHjtoCpQ"
      },
      "outputs": [],
      "source": [
        "tags_df[\"tag\"].value_counts()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1b181a81eb61"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9e7688063ccd"
      },
      "outputs": [],
      "source": [
        "delete_bucket = False\n",
        "\n",
        "# Delete the dataset using the Vertex dataset object\n",
        "dataset.delete()\n",
        "\n",
        "# Undeploy model from the endpoint\n",
        "endpoint.undeploy_all()\n",
        "\n",
        "# Delete the endpoint\n",
        "endpoint.delete()\n",
        "\n",
        "# Delete the model using the Vertex model object\n",
        "model.delete()\n",
        "\n",
        "# Delete the AutoML or Pipeline training job\n",
        "job.delete()\n",
        "\n",
        "# Delete the batch prediction job using the Vertex batch prediction object\n",
        "batch_predict_job.delete()\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "bcKMGrqDrKoS",
        "q3UDGt8MrKoY",
        "bQNL5AjtrKoa",
        "9ewrG6VArKob",
        "eiNL0HSVrKoj",
        "KGT1uT-HrKoj",
        "kyLoUsx9rKok",
        "xksBj3M0rKok",
        "ncP9QwvkrKok",
        "Cyz86ixRYNtT"
      ],
      "name": "SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}