{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h2q27gKz1H20"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TUfAcER1oUS6"
      },
      "outputs": [],
      "source": [
        "#@title 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": "Gb7qyhNL1yWt"
      },
      "source": [
        "# Text classification with TensorFlow Lite Model Maker"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fw5Y7snSuG51"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/lite/tutorials/model_maker_text_classification\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sr3q-gvm3cI8"
      },
      "source": [
        "The TensorFlow Lite Model Maker library simplifies the process of adapting and converting a TensorFlow model to particular input data when deploying this model for on-device ML applications.\n",
        "\n",
        "This notebook shows an end-to-end example that utilizes the Model Maker library to illustrate the adaptation and conversion of a commonly-used text classification model to classify movie reviews on a mobile device. The text classification model classifies text into predefined categories.The inputs should be preprocessed text and the outputs are the probabilities of the categories. The dataset used in this tutorial are positive and negative movie reviews."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcLF2PKkSbV3"
      },
      "source": [
        "## Prerequisites\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2vvAObmTqglq"
      },
      "source": [
        "### Install the required packages\n",
        "To run this example, install the required packages, including the Model Maker package from the [GitHub repo](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "qhl8lqVamEty",
        "outputId": "9757d141-eb26-46e2-fc48-376e0e244142"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting tflite-model-maker\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/13/bc/4c23b9cb9ef612a1f48bac5543bd531665de5eab8f8231111aac067f8c30/tflite_model_maker-0.1.2-py3-none-any.whl (104kB)\n",
            "\r\u001b[K     |███▏                            | 10kB 20.4MB/s eta 0:00:01\r\u001b[K     |██████▎                         | 20kB 5.8MB/s eta 0:00:01\r\u001b[K     |█████████▍                      | 30kB 7.1MB/s eta 0:00:01\r\u001b[K     |████████████▋                   | 40kB 7.7MB/s eta 0:00:01\r\u001b[K     |███████████████▊                | 51kB 6.9MB/s eta 0:00:01\r\u001b[K     |██████████████████▉             | 61kB 7.8MB/s eta 0:00:01\r\u001b[K     |██████████████████████          | 71kB 8.0MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▏      | 81kB 8.4MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▎   | 92kB 7.9MB/s eta 0:00:01\r\u001b[K     |███████████████████████████████▌| 102kB 8.2MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 112kB 8.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: tensorflow-hub\u003e=0.8.0 in /usr/local/lib/python3.6/dist-packages (from tflite-model-maker) (0.9.0)\n",
            "Collecting fire\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/34/a7/0e22e70778aca01a52b9c899d9c145c6396d7b613719cd63db97ffa13f2f/fire-0.3.1.tar.gz (81kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 7.7MB/s \n",
            "\u001b[?25hCollecting flatbuffers==1.12\n",
            "  Downloading https://files.pythonhosted.org/packages/eb/26/712e578c5f14e26ae3314c39a1bdc4eb2ec2f4ddc89b708cf8e0a0d20423/flatbuffers-1.12-py2.py3-none-any.whl\n",
            "Collecting tf-models-nightly\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/d3/e9/c4e5a451c268a5a75a27949562364f6086f6bb33b226a065a8beceefa9ba/tf_models_nightly-2.3.0.dev20200914-py2.py3-none-any.whl (993kB)\n",
            "\u001b[K     |████████████████████████████████| 1.0MB 17.6MB/s \n",
            "\u001b[?25hCollecting sentencepiece\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/d4/a4/d0a884c4300004a78cca907a6ff9a5e9fe4f090f5d95ab341c53d28cbc58/sentencepiece-0.1.91-cp36-cp36m-manylinux1_x86_64.whl (1.1MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1MB 31.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy\u003e=1.17.3 in /usr/local/lib/python3.6/dist-packages (from tflite-model-maker) (1.18.5)\n",
            "Collecting tf-nightly\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/33/d4/61c47ae889b490b9c5f07f4f61bdc057c158a1a1979c375fa019d647a19e/tf_nightly-2.4.0.dev20200914-cp36-cp36m-manylinux2010_x86_64.whl (390.1MB)\n",
            "\u001b[K     |████████████████████████████████| 390.2MB 46kB/s \n",
            "\u001b[?25hRequirement already satisfied: pillow in /usr/local/lib/python3.6/dist-packages (from tflite-model-maker) (7.0.0)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.6/dist-packages (from tflite-model-maker) (0.10.0)\n",
            "Requirement already satisfied: tensorflow-datasets\u003e=2.1.0 in /usr/local/lib/python3.6/dist-packages (from tflite-model-maker) (2.1.0)\n",
            "Collecting tflite-support==0.1.0rc3.dev2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/fa/c5/5e9ee3abd5b4ef8294432cd714407f49a66befa864905b66ee8bdc612795/tflite_support-0.1.0rc3.dev2-cp36-cp36m-manylinux2010_x86_64.whl (1.0MB)\n",
            "\u001b[K     |████████████████████████████████| 1.0MB 50.0MB/s \n",
            "\u001b[?25hRequirement already satisfied: protobuf\u003e=3.8.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-hub\u003e=0.8.0-\u003etflite-model-maker) (3.12.4)\n",
            "Requirement already satisfied: six\u003e=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-hub\u003e=0.8.0-\u003etflite-model-maker) (1.15.0)\n",
            "Requirement already satisfied: termcolor in /usr/local/lib/python3.6/dist-packages (from fire-\u003etflite-model-maker) (1.1.0)\n",
            "Requirement already satisfied: pycocotools in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (2.0.2)\n",
            "Collecting tensorflow-model-optimization\u003e=0.4.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/55/38/4fd48ea1bfcb0b6e36d949025200426fe9c3a8bfae029f0973d85518fa5a/tensorflow_model_optimization-0.5.0-py2.py3-none-any.whl (172kB)\n",
            "\u001b[K     |████████████████████████████████| 174kB 57.7MB/s \n",
            "\u001b[?25hCollecting tf-slim\u003e=1.1.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/02/97/b0f4a64df018ca018cc035d44f2ef08f91e2e8aa67271f6f19633a015ff7/tf_slim-1.1.0-py2.py3-none-any.whl (352kB)\n",
            "\u001b[K     |████████████████████████████████| 358kB 54.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: tensorflow-addons in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (0.8.3)\n",
            "Requirement already satisfied: kaggle\u003e=1.3.9 in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (1.5.8)\n",
            "Requirement already satisfied: oauth2client in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (4.1.3)\n",
            "Collecting seqeval\n",
            "  Downloading https://files.pythonhosted.org/packages/34/91/068aca8d60ce56dd9ba4506850e876aba5e66a6f2f29aa223224b50df0de/seqeval-0.0.12.tar.gz\n",
            "Requirement already satisfied: scipy\u003e=0.19.1 in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (1.4.1)\n",
            "Requirement already satisfied: dataclasses in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (0.7)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (3.2.2)\n",
            "Collecting pyyaml\u003e=5.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/64/c2/b80047c7ac2478f9501676c988a5411ed5572f35d1beff9cae07d321512c/PyYAML-5.3.1.tar.gz (269kB)\n",
            "\u001b[K     |████████████████████████████████| 276kB 55.1MB/s \n",
            "\u001b[?25hRequirement already satisfied: Cython in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (0.29.21)\n",
            "Requirement already satisfied: google-cloud-bigquery\u003e=0.31.0 in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (1.21.0)\n",
            "Collecting opencv-python-headless\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b6/2a/496e06fd289c01dc21b11970be1261c87ce1cc22d5340c14b516160822a7/opencv_python_headless-4.4.0.42-cp36-cp36m-manylinux2014_x86_64.whl (36.6MB)\n",
            "\u001b[K     |████████████████████████████████| 36.6MB 88kB/s \n",
            "\u001b[?25hRequirement already satisfied: psutil\u003e=5.4.3 in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (5.4.8)\n",
            "Requirement already satisfied: gin-config in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (0.3.0)\n",
            "Requirement already satisfied: google-api-python-client\u003e=1.6.7 in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (1.7.12)\n",
            "Requirement already satisfied: pandas\u003e=0.22.0 in /usr/local/lib/python3.6/dist-packages (from tf-models-nightly-\u003etflite-model-maker) (1.0.5)\n",
            "Collecting py-cpuinfo\u003e=3.3.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f6/f5/8e6e85ce2e9f6e05040cf0d4e26f43a4718bcc4bce988b433276d4b1a5c1/py-cpuinfo-7.0.0.tar.gz (95kB)\n",
            "\u001b[K     |████████████████████████████████| 102kB 11.1MB/s \n",
            "\u001b[?25hRequirement already satisfied: grpcio\u003e=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (1.32.0)\n",
            "Requirement already satisfied: opt-einsum\u003e=2.3.2 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (3.3.0)\n",
            "Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (0.3.3)\n",
            "Requirement already satisfied: wheel\u003e=0.26 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (0.35.1)\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (1.6.3)\n",
            "Requirement already satisfied: h5py\u003c2.11.0,\u003e=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (2.10.0)\n",
            "Requirement already satisfied: typing-extensions\u003e=3.7.4.2 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (3.7.4.3)\n",
            "Collecting tb-nightly\u003c3.0.0a0,\u003e=2.4.0a0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/fc/cb/4dfe0d65bffb5e9663261ff664e6f5a2d37672b31dae27a0f14721ac00d3/tb_nightly-2.4.0a20200914-py3-none-any.whl (10.1MB)\n",
            "\u001b[K     |████████████████████████████████| 10.1MB 46.1MB/s \n",
            "\u001b[?25hRequirement already satisfied: google-pasta\u003e=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (0.2.0)\n",
            "Collecting tf-estimator-nightly\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/bd/9a/3bfb9994eda11e426c809ebdf434e2ac5824a0784d980018bb53fd1620ec/tf_estimator_nightly-2.4.0.dev2020091401-py2.py3-none-any.whl (460kB)\n",
            "\u001b[K     |████████████████████████████████| 460kB 51.7MB/s \n",
            "\u001b[?25hRequirement already satisfied: keras-preprocessing\u003c1.2,\u003e=1.1.1 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (1.1.2)\n",
            "Requirement already satisfied: wrapt\u003e=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tf-nightly-\u003etflite-model-maker) (1.12.1)\n",
            "Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (0.24.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (0.16.0)\n",
            "Requirement already satisfied: promise in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (2.3)\n",
            "Requirement already satisfied: requests\u003e=2.19.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (2.23.0)\n",
            "Requirement already satisfied: attrs\u003e=18.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (20.2.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (4.41.1)\n",
            "Requirement already satisfied: dill in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (0.3.2)\n",
            "Collecting pybind11\u003e=2.4\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/89/e3/d576f6f02bc75bacbc3d42494e8f1d063c95617d86648dba243c2cb3963e/pybind11-2.5.0-py2.py3-none-any.whl (296kB)\n",
            "\u001b[K     |████████████████████████████████| 296kB 55.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf\u003e=3.8.0-\u003etensorflow-hub\u003e=0.8.0-\u003etflite-model-maker) (50.3.0)\n",
            "Requirement already satisfied: dm-tree~=0.1.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow-model-optimization\u003e=0.4.1-\u003etf-models-nightly-\u003etflite-model-maker) (0.1.5)\n",
            "Requirement already satisfied: typeguard in /usr/local/lib/python3.6/dist-packages (from tensorflow-addons-\u003etf-models-nightly-\u003etflite-model-maker) (2.7.1)\n",
            "Requirement already satisfied: python-slugify in /usr/local/lib/python3.6/dist-packages (from kaggle\u003e=1.3.9-\u003etf-models-nightly-\u003etflite-model-maker) (4.0.1)\n",
            "Requirement already satisfied: urllib3\u003c1.25,\u003e=1.21.1 in /usr/local/lib/python3.6/dist-packages (from kaggle\u003e=1.3.9-\u003etf-models-nightly-\u003etflite-model-maker) (1.24.3)\n",
            "Requirement already satisfied: certifi in /usr/local/lib/python3.6/dist-packages (from kaggle\u003e=1.3.9-\u003etf-models-nightly-\u003etflite-model-maker) (2020.6.20)\n",
            "Requirement already satisfied: slugify in /usr/local/lib/python3.6/dist-packages (from kaggle\u003e=1.3.9-\u003etf-models-nightly-\u003etflite-model-maker) (0.0.1)\n",
            "Requirement already satisfied: python-dateutil in /usr/local/lib/python3.6/dist-packages (from kaggle\u003e=1.3.9-\u003etf-models-nightly-\u003etflite-model-maker) (2.8.1)\n",
            "Requirement already satisfied: rsa\u003e=3.1.4 in /usr/local/lib/python3.6/dist-packages (from oauth2client-\u003etf-models-nightly-\u003etflite-model-maker) (4.6)\n",
            "Requirement already satisfied: pyasn1\u003e=0.1.7 in /usr/local/lib/python3.6/dist-packages (from oauth2client-\u003etf-models-nightly-\u003etflite-model-maker) (0.4.8)\n",
            "Requirement already satisfied: pyasn1-modules\u003e=0.0.5 in /usr/local/lib/python3.6/dist-packages (from oauth2client-\u003etf-models-nightly-\u003etflite-model-maker) (0.2.8)\n",
            "Requirement already satisfied: httplib2\u003e=0.9.1 in /usr/local/lib/python3.6/dist-packages (from oauth2client-\u003etf-models-nightly-\u003etflite-model-maker) (0.17.4)\n",
            "Requirement already satisfied: Keras\u003e=2.2.4 in /usr/local/lib/python3.6/dist-packages (from seqeval-\u003etf-models-nightly-\u003etflite-model-maker) (2.4.3)\n",
            "Requirement already satisfied: cycler\u003e=0.10 in /usr/local/lib/python3.6/dist-packages (from matplotlib-\u003etf-models-nightly-\u003etflite-model-maker) (0.10.0)\n",
            "Requirement already satisfied: kiwisolver\u003e=1.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib-\u003etf-models-nightly-\u003etflite-model-maker) (1.2.0)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,\u003e=2.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib-\u003etf-models-nightly-\u003etflite-model-maker) (2.4.7)\n",
            "Requirement already satisfied: google-resumable-media!=0.4.0,\u003c0.5.0dev,\u003e=0.3.1 in /usr/local/lib/python3.6/dist-packages (from google-cloud-bigquery\u003e=0.31.0-\u003etf-models-nightly-\u003etflite-model-maker) (0.4.1)\n",
            "Requirement already satisfied: google-cloud-core\u003c2.0dev,\u003e=1.0.3 in /usr/local/lib/python3.6/dist-packages (from google-cloud-bigquery\u003e=0.31.0-\u003etf-models-nightly-\u003etflite-model-maker) (1.0.3)\n",
            "Requirement already satisfied: google-auth-httplib2\u003e=0.0.3 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client\u003e=1.6.7-\u003etf-models-nightly-\u003etflite-model-maker) (0.0.4)\n",
            "Requirement already satisfied: uritemplate\u003c4dev,\u003e=3.0.0 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client\u003e=1.6.7-\u003etf-models-nightly-\u003etflite-model-maker) (3.0.1)\n",
            "Requirement already satisfied: google-auth\u003e=1.4.1 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client\u003e=1.6.7-\u003etf-models-nightly-\u003etflite-model-maker) (1.17.2)\n",
            "Requirement already satisfied: pytz\u003e=2017.2 in /usr/local/lib/python3.6/dist-packages (from pandas\u003e=0.22.0-\u003etf-models-nightly-\u003etflite-model-maker) (2018.9)\n",
            "Requirement already satisfied: markdown\u003e=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (3.2.2)\n",
            "Requirement already satisfied: google-auth-oauthlib\u003c0.5,\u003e=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (0.4.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit\u003e=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (1.7.0)\n",
            "Requirement already satisfied: werkzeug\u003e=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (1.0.1)\n",
            "Requirement already satisfied: googleapis-common-protos\u003c2,\u003e=1.52.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-metadata-\u003etensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (1.52.0)\n",
            "Requirement already satisfied: idna\u003c3,\u003e=2.5 in /usr/local/lib/python3.6/dist-packages (from requests\u003e=2.19.0-\u003etensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (2.10)\n",
            "Requirement already satisfied: chardet\u003c4,\u003e=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests\u003e=2.19.0-\u003etensorflow-datasets\u003e=2.1.0-\u003etflite-model-maker) (3.0.4)\n",
            "Requirement already satisfied: text-unidecode\u003e=1.3 in /usr/local/lib/python3.6/dist-packages (from python-slugify-\u003ekaggle\u003e=1.3.9-\u003etf-models-nightly-\u003etflite-model-maker) (1.3)\n",
            "Requirement already satisfied: google-api-core\u003c2.0.0dev,\u003e=1.14.0 in /usr/local/lib/python3.6/dist-packages (from google-cloud-core\u003c2.0dev,\u003e=1.0.3-\u003egoogle-cloud-bigquery\u003e=0.31.0-\u003etf-models-nightly-\u003etflite-model-maker) (1.16.0)\n",
            "Requirement already satisfied: cachetools\u003c5.0,\u003e=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth\u003e=1.4.1-\u003egoogle-api-python-client\u003e=1.6.7-\u003etf-models-nightly-\u003etflite-model-maker) (4.1.1)\n",
            "Requirement already satisfied: importlib-metadata; python_version \u003c \"3.8\" in /usr/local/lib/python3.6/dist-packages (from markdown\u003e=2.6.8-\u003etb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (1.7.0)\n",
            "Requirement already satisfied: requests-oauthlib\u003e=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib\u003c0.5,\u003e=0.4.1-\u003etb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (1.3.0)\n",
            "Requirement already satisfied: zipp\u003e=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata; python_version \u003c \"3.8\"-\u003emarkdown\u003e=2.6.8-\u003etb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (3.1.0)\n",
            "Requirement already satisfied: oauthlib\u003e=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib\u003e=0.7.0-\u003egoogle-auth-oauthlib\u003c0.5,\u003e=0.4.1-\u003etb-nightly\u003c3.0.0a0,\u003e=2.4.0a0-\u003etf-nightly-\u003etflite-model-maker) (3.1.0)\n",
            "Building wheels for collected packages: fire, seqeval, pyyaml, py-cpuinfo\n",
            "  Building wheel for fire (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for fire: filename=fire-0.3.1-py2.py3-none-any.whl size=111005 sha256=9eaa2d36e17621d136f8ab1707a5a4e8994c53d5076a9edde21aab7696ba3e09\n",
            "  Stored in directory: /root/.cache/pip/wheels/c1/61/df/768b03527bf006b546dce284eb4249b185669e65afc5fbb2ac\n",
            "  Building wheel for seqeval (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for seqeval: filename=seqeval-0.0.12-cp36-none-any.whl size=7423 sha256=1ce4604da2a395f0304db708bf2e2c1831033ed8b1f7c23927d70ed9ed7b7110\n",
            "  Stored in directory: /root/.cache/pip/wheels/4f/32/0a/df3b340a82583566975377d65e724895b3fad101a3fb729f68\n",
            "  Building wheel for pyyaml (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyyaml: filename=PyYAML-5.3.1-cp36-cp36m-linux_x86_64.whl size=44619 sha256=d51b6ef3e90de74d0c1cee8f7aafe0a6d8674348c8437cd89ad5c60a6c3dc726\n",
            "  Stored in directory: /root/.cache/pip/wheels/a7/c1/ea/cf5bd31012e735dc1dfea3131a2d5eae7978b251083d6247bd\n",
            "  Building wheel for py-cpuinfo (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for py-cpuinfo: filename=py_cpuinfo-7.0.0-cp36-none-any.whl size=20071 sha256=096439bff3cb3e4cc21b86472c629017fd9c972d6e2ed231e1a91d2096fc687d\n",
            "  Stored in directory: /root/.cache/pip/wheels/f1/93/7b/127daf0c3a5a49feb2fecd468d508067c733fba5192f726ad1\n",
            "Successfully built fire seqeval pyyaml py-cpuinfo\n",
            "Installing collected packages: fire, flatbuffers, tensorflow-model-optimization, tf-slim, seqeval, pyyaml, opencv-python-headless, sentencepiece, tb-nightly, tf-estimator-nightly, tf-nightly, py-cpuinfo, tf-models-nightly, pybind11, tflite-support, tflite-model-maker\n",
            "  Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "Successfully installed fire-0.3.1 flatbuffers-1.12 opencv-python-headless-4.4.0.42 py-cpuinfo-7.0.0 pybind11-2.5.0 pyyaml-5.3.1 sentencepiece-0.1.91 seqeval-0.0.12 tb-nightly-2.4.0a20200914 tensorflow-model-optimization-0.5.0 tf-estimator-nightly-2.4.0.dev2020091401 tf-models-nightly-2.3.0.dev20200914 tf-nightly-2.4.0.dev20200914 tf-slim-1.1.0 tflite-model-maker-0.1.2 tflite-support-0.1.0rc3.dev2\n"
          ]
        }
      ],
      "source": [
        "!pip install tflite-model-maker"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l6lRhVK9Q_0U"
      },
      "source": [
        "Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtxiUeZEiXpt"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import os\n",
        "\n",
        "import tensorflow as tf\n",
        "assert tf.__version__.startswith('2')\n",
        "\n",
        "from tflite_model_maker import configs\n",
        "from tflite_model_maker import ExportFormat\n",
        "from tflite_model_maker import model_spec\n",
        "from tflite_model_maker import text_classifier\n",
        "from tflite_model_maker import TextClassifierDataLoader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BRd13bfetO7B"
      },
      "source": [
        "### Get the data path\n",
        "Download the dataset for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "id": "R2BSkxWg6Rhx",
        "outputId": "972b735a-d5f6-4152-9c0f-12b7b97b8d86"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading data from https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSST-2.zip?alt=media\u0026token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8\n",
            "7446528/7439277 [==============================] - 0s 0us/step\n"
          ]
        }
      ],
      "source": [
        "data_dir = tf.keras.utils.get_file(\n",
        "      fname='SST-2.zip',\n",
        "      origin='https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSST-2.zip?alt=media\u0026token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8',\n",
        "      extract=True)\n",
        "data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6MSCjPAvs2EQ"
      },
      "source": [
        "You can also upload your own dataset to work through this tutorial. Upload your dataset by using the left sidebar in Colab.\n",
        "\n",
        "\u003cimg src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_text_classification.png\" alt=\"Upload File\" width=\"800\" hspace=\"100\"\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uO5egTlrtWxm"
      },
      "source": [
        "If you prefer not to upload your dataset to the cloud, you can also locally run the library by following the [guide](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xushUyZXqP59"
      },
      "source": [
        "## End-to-End Workflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WlKU3SMX6TnB"
      },
      "source": [
        "This workflow consists of five steps as outlined below:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PBPUIhEjMjTR"
      },
      "source": [
        "Step 1. Choose a model specification that represents a text classification model.\n",
        "\n",
        "This tutorial uses [MobileBERT](https://arxiv.org/pdf/2004.02984.pdf) as an example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CtdZ-JDwMimd"
      },
      "outputs": [],
      "source": [
        "spec = model_spec.get('mobilebert_classifier')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s5U-A3tw6Y27"
      },
      "source": [
        "Step 2.   Load train and test data specific to an on-device ML app and preprocess the data according to a specific `model_spec`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HD5BvzWe6YKa"
      },
      "outputs": [],
      "source": [
        "train_data = TextClassifierDataLoader.from_csv(\n",
        "      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),\n",
        "      text_column='sentence',\n",
        "      label_column='label',\n",
        "      model_spec=spec,\n",
        "      delimiter='\\t',\n",
        "      is_training=True)\n",
        "test_data = TextClassifierDataLoader.from_csv(\n",
        "      filename=os.path.join(os.path.join(data_dir, 'dev.tsv')),\n",
        "      text_column='sentence',\n",
        "      label_column='label',\n",
        "      model_spec=spec,\n",
        "      delimiter='\\t',\n",
        "      is_training=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2uZkLR6N6gDR"
      },
      "source": [
        "Step 3. Customize the TensorFlow model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kwlYdTcg63xy"
      },
      "outputs": [],
      "source": [
        "model = text_classifier.create(train_data, model_spec=spec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-BzCHLWJ6h7q"
      },
      "source": [
        "Step 4. Evaluate the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8xmnl6Yy7ARn"
      },
      "outputs": [],
      "source": [
        "loss, acc = model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CgCDMe0e6jlT"
      },
      "source": [
        "Step 5.  Export as a TensorFlow Lite model with [metadata](https://www.tensorflow.org/lite/convert/metadata).\n",
        "\n",
        "Since MobileBERT is too big for on-device applications, use [dynamic range quantization](https://www.tensorflow.org/lite/performance/post_training_quantization#dynamic_range_quantization) on the model to compress it by almost 4x with minimal performance degradation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZQRLmkGumr9Y"
      },
      "outputs": [],
      "source": [
        "config = configs.QuantizationConfig.create_dynamic_range_quantization(optimizations=[tf.lite.Optimize.OPTIMIZE_FOR_LATENCY])\n",
        "config._experimental_new_quantizer = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hm_UULdW7A9T"
      },
      "outputs": [],
      "source": [
        "model.export(export_dir='mobilebert/', quantization_config=config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rVxaf3x_7OfB"
      },
      "source": [
        "You can also download the model using the left sidebar in Colab.\n",
        "\n",
        "After executing the 5 steps above, you can further use the TensorFlow Lite model file in on-device applications using [BertNLClassifier API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/bert_nl_classifier) in [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l65ctmtW7_FF"
      },
      "source": [
        "The following sections walk through the example step by step to show more detail."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kJ_B8fMDOhMR"
      },
      "source": [
        "## Choose a `model_spec` that Represents a Model for Text Classifier\n",
        "\n",
        "Each `model_spec` object represents a specific model for the text classifier. TensorFlow Lite Model Maker currently supports [MobileBERT](https://arxiv.org/pdf/2004.02984.pdf), averaging word embeddings and [BERT-Base](https://arxiv.org/pdf/1810.04805.pdf) models.\n",
        "\n",
        "Supported Model | Name of model_spec | Model Description\n",
        "--- | --- | ---\n",
        "MobileBERT | 'mobilebert_classifier' | 4.3x smaller and 5.5x faster than BERT-Base while achieving competitive results, suitable for on-device applications.\n",
        "BERT-Base | 'bert_classifier' | Standard BERT model that is widely used in NLP tasks.\n",
        "averaging word embedding | 'average_word_vec' | Averaging text word embeddings with RELU activation.\n",
        "\n",
        "This tutorial uses a smaller model, `average_word_vec` that you can retrain multiple times to demonstrate the process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vEAWuZQ1PFiX"
      },
      "outputs": [],
      "source": [
        "spec = model_spec.get('average_word_vec')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ygEncJxtl-nQ"
      },
      "source": [
        "## Load Input Data Specific to an On-device ML App\n",
        "\n",
        "The [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank) is one of the tasks in the [GLUE](https://gluebenchmark.com/) benchmark. It contains 67,349 movie reviews for training and 872 movie reviews for validation. The dataset has two classes: positive and negative movie reviews.\n",
        "\n",
        "Download the archived version of the dataset and extract it.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7tOfUr2KlgpU"
      },
      "outputs": [],
      "source": [
        "data_dir = tf.keras.utils.get_file(\n",
        "      fname='SST-2.zip',\n",
        "      origin='https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSST-2.zip?alt=media\u0026token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8',\n",
        "      extract=True)\n",
        "data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E051HBUM5owi"
      },
      "source": [
        "The SST-2 dataset has `train.tsv` for training and `dev.tsv` for validation. The files have the following format:\n",
        "\n",
        "sentence | label\n",
        "--- | ---\n",
        "it 's a charming and often affecting journey . | 1\n",
        "unflinchingly bleak and desperate | 0\n",
        "\n",
        "A positive review is labeled 1 and a negative review is labeled 0.\n",
        "\n",
        "Use the `TestClassifierDataLoader.from_csv` method to load the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I_fOlZsklmlL"
      },
      "outputs": [],
      "source": [
        "train_data = TextClassifierDataLoader.from_csv(\n",
        "      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),\n",
        "      text_column='sentence',\n",
        "      label_column='label',\n",
        "      model_spec=spec,\n",
        "      delimiter='\\t',\n",
        "      is_training=True)\n",
        "test_data = TextClassifierDataLoader.from_csv(\n",
        "      filename=os.path.join(os.path.join(data_dir, 'dev.tsv')),\n",
        "      text_column='sentence',\n",
        "      label_column='label',\n",
        "      model_spec=spec,\n",
        "      delimiter='\\t',\n",
        "      is_training=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MlHvVvv2hw4H"
      },
      "source": [
        "The Model Maker library also supports the `from_folder()` method to load data. It assumes that the text data of the same class are in the same subdirectory and that the subfolder name is the class name. Each text file contains one movie review sample. The `class_labels` parameter is used to specify which the subfolders."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AWuoensX4vDA"
      },
      "source": [
        "## Customize the TensorFlow Model\n",
        "\n",
        "Create a custom text classifier model based on the loaded data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TvYSUuJY3QxR"
      },
      "outputs": [],
      "source": [
        "model = text_classifier.create(train_data, model_spec=spec, epochs=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0JKI-pNc8idH"
      },
      "source": [
        "Examine the detailed model structure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gd7Hs8TF8n3H"
      },
      "outputs": [],
      "source": [
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LP5FPk_tOxoZ"
      },
      "source": [
        "## Evaluate the Customized Model\n",
        "\n",
        "Evaluate the model with the test data and get its loss and accuracy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A8c2ZQ0J3Riy"
      },
      "outputs": [],
      "source": [
        "loss, acc = model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aeHoGAceO2xV"
      },
      "source": [
        "## Export as a TensorFlow Lite Model\n",
        "\n",
        "Convert the existing model to TensorFlow Lite model format with [metadata](https://www.tensorflow.org/lite/convert/metadata) that you can later use in an on-device ML application. The label file and the vocab file are embedded in metadata. The default TFLite filename is `model.tflite`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Im6wA9lK3TQB"
      },
      "outputs": [],
      "source": [
        "model.export(export_dir='average_word_vec/')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w12kvDdHJIGH"
      },
      "source": [
        "The TensorFlow Lite model file can be used in the [text classification](https://github.com/tensorflow/examples/tree/master/lite/examples/text_classification) reference app using [NLClassifier API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/nl_classifier) in [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AVy0ormoMZwL"
      },
      "source": [
        "The allowed export formats can be one or a list of the following:\n",
        "\n",
        "*   `ExportFormat.TFLITE`\n",
        "*   `ExportFormat.LABEL`\n",
        "*   `ExportFormat.VOCAB`\n",
        "*   `ExportFormat.SAVED_MODEL`\n",
        "\n",
        "By default, it just exports TensorFlow Lite model with metadata. You can also selectively export different files. For instance, exporting only the label file and vocab file as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nbK7nzK_Mfx4"
      },
      "outputs": [],
      "source": [
        "model.export(export_dir='average_word_vec/', export_format=[ExportFormat.LABEL, ExportFormat.VOCAB])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HZKYthlVrTos"
      },
      "source": [
        "You can evaluate the tflite model with `evaluate_tflite` method to get its accuracy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ochbq95ZrVFX"
      },
      "outputs": [],
      "source": [
        "accuracy = model.evaluate_tflite('average_word_vec/model.tflite', test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EoWiA_zX8rxE"
      },
      "source": [
        "## Advanced Usage\n",
        "\n",
        "The `create` function is the driver function that the Model Maker library uses to create models. The `model_spec` parameter defines the model specification. The `AverageWordVecModelSpec` and `BertClassifierModelSpec` classes are currently supported. The `create` function comprises of the following steps:\n",
        "\n",
        "1. Creates the model for the text classifier according to `model_spec`.\n",
        "2. Trains the classifier model.  The default epochs and the default batch size are set by the `default_training_epochs` and `default_batch_size` variables in the `model_spec` object.\n",
        "\n",
        "This section covers advanced usage topics like adjusting the model and the training hyperparameters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mwtiksguDfhl"
      },
      "source": [
        "### Adjust the model\n",
        "\n",
        "You can adjust the model infrastructure like the `wordvec_dim` and the `seq_len` variables in the `AverageWordVecModelSpec` class.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cAOd5_bzH9AQ"
      },
      "source": [
        "For example, you can train the model with a larger value of `wordvec_dim`. Note that you must construct a new `model_spec` if you modify the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e9WBN0UTQoMN"
      },
      "outputs": [],
      "source": [
        "new_model_spec = model_spec.AverageWordVecModelSpec(wordvec_dim=32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LSTdghTP0Cv"
      },
      "source": [
        "Get the preprocessed data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DVZurFBORG3J"
      },
      "outputs": [],
      "source": [
        "new_train_data = TextClassifierDataLoader.from_csv(\n",
        "      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),\n",
        "      text_column='sentence',\n",
        "      label_column='label',\n",
        "      model_spec=new_model_spec,\n",
        "      delimiter='\\t',\n",
        "      is_training=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tD7QVVHeRZoM"
      },
      "source": [
        "Train the new model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PzpV246_JGEu"
      },
      "outputs": [],
      "source": [
        "model = text_classifier.create(new_train_data, model_spec=new_model_spec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E8VxPiOLy4Gv"
      },
      "source": [
        "You can also adjust the MobileBERT model.\n",
        "\n",
        "The model parameters you can adjust are:\n",
        "\n",
        "* `seq_len`: Length of the sequence to feed into the model.\n",
        "* `initializer_range`: The standard deviation of the `truncated_normal_initializer` for initializing all weight matrices.\n",
        "* `trainable`: Boolean that specifies whether the pre-trained layer is trainable.\n",
        "\n",
        "The training pipeline parameters you can adjust are:\n",
        "\n",
        "* `model_dir`: The location of the model checkpoint files. If not set, a temporary directory will be used.\n",
        "* `dropout_rate`: The dropout rate.\n",
        "* `learning_rate`: The initial learning rate for the Adam optimizer.\n",
        "* `tpu`: TPU address to connect to.\n",
        "\n",
        "For instance, you can set the `seq_len=256` (default is 128). This allows the model to classify longer text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4tr9BLcjy4Sh"
      },
      "outputs": [],
      "source": [
        "new_model_spec = model_spec.get('mobilebert_classifier')\n",
        "new_model_spec.seq_len = 256"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LvQuy7RSDir3"
      },
      "source": [
        "### Tune the training hyperparameters\n",
        "You can also tune the training hyperparameters like `epochs` and `batch_size` that affect the model accuracy. For instance,\n",
        "\n",
        "*   `epochs`: more epochs could achieve better accuracy, but may lead to overfitting.\n",
        "*   `batch_size`: the number of samples to use in one training step.\n",
        "\n",
        "For example, you can train with more epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rnWFaYZBG6NW"
      },
      "outputs": [],
      "source": [
        "model = text_classifier.create(train_data, model_spec=spec, epochs=20)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nUaKQZBQHBQR"
      },
      "source": [
        "Evaluate the newly retrained model with 20 training epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BMPi1xflHDSY"
      },
      "outputs": [],
      "source": [
        "loss, accuracy = model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eq6B9lKMfhS6"
      },
      "source": [
        "### Change the Model Architecture\n",
        "\n",
        "You can change the model by changing the `model_spec`. The following shows how to change to BERT-Base model.\n",
        "\n",
        "Change the `model_spec` to BERT-Base model for the text classifier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QfFCWrwyggrT"
      },
      "outputs": [],
      "source": [
        "spec = model_spec.get('bert_classifier')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2d7yycrgu6L"
      },
      "source": [
        "The remaining steps are the same."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "model_maker_text_classification.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
