notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": "TirJ-SGQseby"
      },
      "source": [
        "# Vertex AI Model Garden MediaPipe with text classification\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.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",
        "\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.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/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb\">\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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dwGLvtIeECLK"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9\n",
        "\n",
        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to train an on-device text classification model in Vertex AI Model Garden.\n",
        "\n",
        "### Objective\n",
        "\n",
        "* Train new models\n",
        "  * Convert input data to training formats\n",
        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new models\n",
        "  * Export models\n",
        "\n",
        "* Cleanup resources\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z__i0w0lCAsW"
      },
      "source": [
        "### Colab only\n",
        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade pip\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)\n",
        "\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tTy1gX11kCJY"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
        "assert REGION_PREFIX in (\n",
        "    \"us\",\n",
        "    \"europe\",\n",
        "    \"asia\",\n",
        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wExiMUxFk91"
      },
      "outputs": [],
      "source": [
        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
        "\n",
        "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
        "EVALUATION_RESULT_OUTPUT_FILE = os.path.join(\n",
        "    EVALUATION_RESULT_OUTPUT_DIRECTORY, \"evaluation.json\"\n",
        ")\n",
        "\n",
        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"model.tflite\"\n",
        ")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n6IFz75WGCam"
      },
      "source": [
        "### Define training machine specs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_text_classifier_%s\" % now\n",
        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "TRAINING_ACCELERATOR_COUNT = 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-rsdAcBV-vlf"
      },
      "source": [
        "## Train your customized models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Get the Dataset\n",
        "\n",
        "The following code block uses the [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank) dataset which contains 67,349 movie reviews for training and 872 movie reviews for testing. The dataset has two classes: positive and negative movie reviews. Positive reviews are labeled with 1 and negative reviews with 0.\n",
        "\n",
        "The SST-2 dataset is stored as a TSV file. The only difference between the TSV and CSV formats is that TSV uses a tab `\\t` character as its delimiter and CSV uses a comma `,`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IndQ_m6ddUEM"
      },
      "outputs": [],
      "source": [
        "training_data_path = (\n",
        "    \"gs://mediapipe-tasks/text_classifier/SST-2/train.tsv\"  # @param {type:\"string\"}\n",
        ")\n",
        "validation_data_path = (\n",
        "    \"gs://mediapipe-tasks/text_classifier/SST-2/dev.tsv\"  # @param {type:\"string\"}\n",
        ")\n",
        "\n",
        "# The delimiter used in the dataset.\n",
        "delimiter = \"\\t\"  # @param {type:\"string\"}\n",
        "\n",
        "# Character used to quote fields that contain special characters\n",
        "# like the `delimiter`.\n",
        "quotechar = \"\\t\"  # @param {type:\"string\"}\n",
        "\n",
        "# Sequence of keys for the CSV columns (represented as a comma\n",
        "# separated list). If empty, the first row of the CSV file is used\n",
        "# as the keys\n",
        "fieldnames = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Column name for the input text.\n",
        "text_column = \"sentence\"  # @param {type:\"string\"}\n",
        "\n",
        "# Column name for the labels.\n",
        "label_column = \"label\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaff6f5be7f6"
      },
      "source": [
        "### Set fine-tuning options\n",
        "\n",
        "You can pick between different model architectures to further customize your training:\n",
        "\n",
        "*   Average Word Embedding Model\n",
        "*   BERT-classifier\n",
        "\n",
        "To set the model architecture and other training parameters, adjust the following values:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "um_XKbmpTaHx"
      },
      "outputs": [],
      "source": [
        "model_architecture = (\n",
        "    \"average_word_embedding\"  # @param [\"average_word_embedding\", \"mobilebert\"]\n",
        ")\n",
        "\n",
        "# The learning rate to use for gradient descent-based\n",
        "# optimizers. Defaults to 3e-5 for the BERT-based classifier\n",
        "# and 0 for the average word-embedding classifier because\n",
        "# it does not need such an optimizer.\n",
        "learning_rate: float = 0.0  # @param {type:\"number\"}\n",
        "\n",
        "# Batch size for training. Defaults to 32 for the average\n",
        "# word-embedding classifier and 48 for the BERT-based\n",
        "# classifier.\n",
        "batch_size: int = 48  # @param {type:\"number\"}\n",
        "\n",
        "# Number of training iterations over the dataset. Defaults\n",
        "# to 10 for the average word-embedding classifier and 3\n",
        "# for the BERT-based classifier.\n",
        "epochs: int = 10  # @param {type:\"slider\", min:0, max:100, step:1}\n",
        "\n",
        "# An integer that indicates the number of training steps per\n",
        "# epoch. If set to 0, the training pipeline calculates the\n",
        "# default steps per epoch as the training dataset size\n",
        "# divided by batch size.\n",
        "steps_per_epoch: int = 0  # @param {type:\"number\"}\n",
        "\n",
        "# Controls whether the dataset is shuffled before training.\n",
        "shuffle: bool = False  # @param {type:\"boolean\"}\n",
        "\n",
        "# Length of the sequence to feed into the model.\n",
        "seq_len: int = 256  # @param {type:\"number\"}\n",
        "\n",
        "# Whether to convert all uppercase characters to lowercase\n",
        "# during preprocessing.\n",
        "do_lower_case: bool = True  # @param {type:\"boolean\"}\n",
        "\n",
        "# The rate for dropout.\n",
        "dropout_rate: float = 0.2  # @param {type:\"number\"}\n",
        "\n",
        "# Dimension of the word embedding. Only used for the Average Word\n",
        "# Embedding Model.\n",
        "wordvec_dim: int = 16  # @param {type:\"number\"}\n",
        "\n",
        "# Number of words to generate the vocabulary from data.\n",
        "# Only used for the Average Word Embedding Model.\n",
        "vocab_size: int = 10000  # @param {type:\"number\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HwcCjwlBTQIz"
      },
      "source": [
        "### Run fine-tuning\n",
        "With your training dataset and fine-tuning options prepared, you are ready to start the fine-tuning process. This process is resource intensive and can take a few minutes to a few hours depending on the model archtiecture and your available compute resources. On Vertex AI with GPU processing, the example fine-tuning below takes between 2-3 minutes to train an Average Word Embedding Model on the SST-2 dataset.\n",
        "\n",
        "To begin the fine-tuning process, use the following code:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
        "evaluation_result_path = EVALUATION_RESULT_OUTPUT_DIRECTORY\n",
        "\n",
        "preprocessing_params = {\n",
        "    \"text_column\": text_column,\n",
        "    \"label_column\": label_column,\n",
        "    \"delimiter\": delimiter,\n",
        "    \"quotechar\": quotechar,\n",
        "}\n",
        "if fieldnames:\n",
        "    preprocessing_params[\"fieldnames\"] = [\n",
        "        fieldname.strip() for fieldname in fieldnames.split(\",\")\n",
        "    ]\n",
        "\n",
        "hparams = {\n",
        "    \"learning_rate\": learning_rate,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"epochs\": epochs,\n",
        "    \"shuffle\": shuffle,\n",
        "}\n",
        "if steps_per_epoch:\n",
        "    hparams[\"steps_per_epoch\"] = steps_per_epoch\n",
        "\n",
        "model_options = {\n",
        "    \"dropout_rate\": dropout_rate,\n",
        "    \"wordvec_dim\": wordvec_dim,\n",
        "    \"do_lower_case\": do_lower_case,\n",
        "    \"vocab_size\": vocab_size,\n",
        "    \"dropout_rate\": dropout_rate,\n",
        "}\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAINING_CONTAINER,\n",
        "            \"command\": [],\n",
        "            \"args\": [\n",
        "                \"--task_name=text_classifier\",\n",
        "                \"--training_data_path=%s\" % training_data_path,\n",
        "                \"--validation_data_path=%s\" % validation_data_path,\n",
        "                \"--evaluation_result_path=%s\" % evaluation_result_path,\n",
        "                \"--model_export_path=%s\" % model_export_path,\n",
        "                \"--model_architecture=%s\" % model_architecture,\n",
        "                \"--preprocessing_params=%s\" % json.dumps(preprocessing_params),\n",
        "                \"--hparams=%s\" % json.dumps(hparams),\n",
        "                \"--model_options=%s\" % json.dumps(model_options),\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "training_job = aiplatform.CustomJob(\n",
        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "training_job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rXMF2tnV_WS0"
      },
      "source": [
        "## Export model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g0BGaofgsMsy"
      },
      "source": [
        "After finetuning, you can save the Tensorflow Lite model, try it out in the [Text Classification](https://mediapipe-studio.webapps.google.com/demo/text_classifier) demo in MediaPipe Studio or integrate it with your on-device application by following the [Text classification task guide](https://developers.google.com/mediapipe/solutions/text/text_classifier). The exported model contains the generates required model metadata, as well as a classification label file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "\n",
        "def copy_model(model_source, model_dest):\n",
        "    ! gsutil cp {model_source} {model_dest}\n",
        "\n",
        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"text_classification_model.tflite\")\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import files\n",
        "\n",
        "    files.download(\"text_classification_model.tflite\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "# Delete training data and jobs.\n",
        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
        "    training_job.delete()\n",
        "\n",
        "!gsutil rm -r {STAGING_BUCKET}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_mediapipe_text_classification.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}