notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_mediapipe_text_classification.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "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 pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
        "\n",
        "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
        "\n",
        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
        "# @markdown | ----------- | ----------- | ----------- |\n",
        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
        "\n",
        "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import json\n",
        "import os\n",
        "import uuid\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "if not REGION:\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"mediapipe_text_classification\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
        "\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": "-rsdAcBV-vlf"
      },
      "source": [
        "## Train your customized models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "IndQ_m6ddUEM"
      },
      "outputs": [],
      "source": [
        "# @title Set the dataset\n",
        "\n",
        "# @markdown 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",
        "# @markdown 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",
        "\n",
        "training_data_path = \"gs://mediapipe-tasks/text_classifier/SST-2/train.tsv\"  # @param {type:\"string\"}\n",
        "\n",
        "validation_data_path = \"gs://mediapipe-tasks/text_classifier/SST-2/dev.tsv\"  # @param {type:\"string\"}\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": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "um_XKbmpTaHx"
      },
      "outputs": [],
      "source": [
        "# @title Set fine-tuning options\n",
        "\n",
        "# @markdown You can pick between different model architectures to further customize your training:\n",
        "# @markdown *   Average Word Embedding Model\n",
        "# @markdown *   BERT-classifier\n",
        "\n",
        "# @markdown To set the model architecture and other training parameters, adjust the below values:\n",
        "\n",
        "model_architecture = \"average_word_embedding\"  # @param [\"average_word_embedding\", \"mobilebert\"]\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": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "# @title Run finetuning job\n",
        "\n",
        "# @markdown 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",
        "# @markdown To begin the fine-tuning process, use the following code:\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",
        "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",
        "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\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",
        "# Check quota.\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=TRAINING_ACCELERATOR_TYPE,\n",
        "    accelerator_count=2,\n",
        "    is_for_training=True,\n",
        ")\n",
        "\n",
        "\n",
        "# Add labels for the finetuning job.\n",
        "labels = {\n",
        "    \"mg-source\": \"notebook\",\n",
        "    \"mg-notebook-name\": \"model_garden_mediapipe_text_classification.ipynb\".split(\".\")[\n",
        "        0\n",
        "    ],\n",
        "}\n",
        "\n",
        "labels[\"mg-tune\"] = \"publishers-google-models-mdeiapipe\"\n",
        "versioned_model_id = model_architecture.lower().replace(\"_\", \"-\")\n",
        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\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",
        "    labels=labels,\n",
        ")\n",
        "\n",
        "training_job.run()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "# @title Export model\n",
        "\n",
        "# @markdown 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.\n",
        "\n",
        "! gsutil cp $EXPORTED_MODEL_OUTPUT_FILE text_classification_model.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "# @title Clean up training jobs and buckets\n",
        "# @markdown Delete temporary GCS buckets.\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME\n",
        "\n",
        "# Delete training data and jobs.\n",
        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
        "    training_job.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_mediapipe_text_classification.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}