Improve Model Maker tutorial notebook

PiperOrigin-RevId: 358338694
Change-Id: I2f6538530d2cacc9ee56ceba8a662df830e145d2
diff --git a/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb b/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb
index e9c71a5..c8604e2 100644
--- a/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb
+++ b/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb
@@ -46,20 +46,20 @@
         "id": "Fw5Y7snSuG51"
       },
       "source": [
-        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
-        "  <td>\n",
-        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/tutorials/model_maker_text_classification\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
-        "  </td>\n",
-        "</table>"
+        "\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"
       ]
     },
     {
@@ -89,7 +89,9 @@
       },
       "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)."
+        "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).\n",
+        "\n",
+        "**If you run this notebook on Colab, you may see an error message about `tensorflowjs` and `tensorflow-hub` version imcompatibility. It is safe to ignore this error as we do not use `tensorflowjs` in this workflow.**"
       ]
     },
     {
@@ -100,7 +102,7 @@
       },
       "outputs": [],
       "source": [
-        "!pip install tflite-model-maker"
+        "!pip install -q tflite-model-maker"
       ]
     },
     {
@@ -123,14 +125,15 @@
         "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"
+        "from tflite_model_maker import TextClassifierDataLoader\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "assert tf.__version__.startswith('2')\n",
+        "tf.get_logger().setLevel('ERROR')"
       ]
     },
     {
@@ -139,8 +142,9 @@
         "id": "BRd13bfetO7B"
       },
       "source": [
-        "### Get the data path\n",
-        "Download the dataset for this tutorial."
+        "### Download the sample training data.\n",
+        "\n",
+        "In this tutorial, we will use the [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank) which 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 testing. The dataset has two classes: positive and negative movie reviews."
       ]
     },
     {
@@ -161,21 +165,53 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "6MSCjPAvs2EQ"
+        "id": "gPYTbGrizcTC"
       },
       "source": [
-        "You can also upload your own dataset to work through this tutorial. Upload your dataset by using the left sidebar in Colab.\n",
+        "The SST-2 dataset is stored in TSV format. The only difference between TSV and CSV is that TSV uses a tab `\\t` character as its delimiter instead of a comma `,` in the CSV format.\n",
         "\n",
-        "<img src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_text_classification.png\" alt=\"Upload File\" width=\"800\" hspace=\"100\">\n"
+        "Here are the first 5 lines of the training dataset. label=0 means negative, label=1 means positive.\n",
+        "\n",
+        "| sentence                                                                                  | label |   |   |   |\n",
+        "|-------------------------------------------------------------------------------------------|-------|---|---|---|\n",
+        "| hide new secretions from the parental units                                               | 0     |   |   |   |\n",
+        "| contains no wit , only labored gags                                                       | 0     |   |   |   |\n",
+        "| that loves its characters and communicates something rather beautiful about human nature  | 1     |   |   |   |\n",
+        "| remains utterly satisfied to remain the same throughout                                   | 0     |   |   |   |\n",
+        "| on the worst revenge-of-the-nerds clichés the filmmakers could dredge up                  | 0     |   |   |   |\n",
+        "\n",
+        "Next, we will load the dataset into a Pandas dataframe and change the current label names (`0` and `1`) to a more human-readable ones (`negative` and `positive`) and use them for model training.\n",
+        "\n"
       ]
     },
     {
-      "cell_type": "markdown",
+      "cell_type": "code",
+      "execution_count": null,
       "metadata": {
-        "id": "uO5egTlrtWxm"
+        "id": "iLNaOXnl3JQB"
       },
+      "outputs": [],
       "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)."
+        "import pandas as pd\n",
+        "\n",
+        "def replace_label(original_file, new_file):\n",
+        "  # Load the original file to pandas. We need to specify the separator as\n",
+        "  # '\\t' as the training data is stored in TSV format\n",
+        "  df = pd.read_csv(original_file, sep='\\t')\n",
+        "\n",
+        "  # Define how we want to change the label name\n",
+        "  label_map = {0: 'negative', 1: 'positive'}\n",
+        "\n",
+        "  # Excute the label change\n",
+        "  df.replace({'label': label_map}, inplace=True)\n",
+        "\n",
+        "  # Write the updated dataset to a new file\n",
+        "  df.to_csv(new_file)\n",
+        "\n",
+        "# Replace the label name for both the training and test dataset. Then write the\n",
+        "# updated CSV dataset to the current folder.\n",
+        "replace_label(os.path.join(os.path.join(data_dir, 'train.tsv')), 'train.csv')\n",
+        "replace_label(os.path.join(os.path.join(data_dir, 'dev.tsv')), 'dev.csv')"
       ]
     },
     {
@@ -184,27 +220,13 @@
         "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",
+        "## Quickstart\n",
         "\n",
-        "This tutorial uses [MobileBERT](https://arxiv.org/pdf/2004.02984.pdf) as an example."
+        "There are five steps to train a text classification model:\n",
+        "\n",
+        "**Step 1. Choose a text classification model archiecture.**\n",
+        "\n",
+        "Here we use the average word embedding model architecture, which will produce a small and fast model with decent accuracy."
       ]
     },
     {
@@ -215,7 +237,16 @@
       },
       "outputs": [],
       "source": [
-        "spec = model_spec.get('mobilebert_classifier')"
+        "spec = model_spec.get('average_word_vec')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "yug6gR9qyHui"
+      },
+      "source": [
+        "Model Maker also supports other model architectures such as [BERT](https://arxiv.org/abs/1810.04805). If you are interested to learn about other architecture, see the [Choose a model architecture for Text Classifier](#scrollTo=kJ_B8fMDOhMR) section below."
       ]
     },
     {
@@ -224,7 +255,11 @@
         "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`."
+        "**Step 2.   Load the training and test data, then preprocess them according to a specific `model_spec`.**\n",
+        "\n",
+        "Model Maker can take input data in the CSV format. We will load the training and test dataset with the human-readable label name that were created earlier.\n",
+        "\n",
+        "Each model architecture requires input data to be processed in a particular way. `TextClassifierDataLoader` reads the requirement from `model_spec` and automatically execute the necessary preprocessing."
       ]
     },
     {
@@ -236,18 +271,16 @@
       "outputs": [],
       "source": [
         "train_data = TextClassifierDataLoader.from_csv(\n",
-        "      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),\n",
+        "      filename='train.csv',\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",
+        "      filename='dev.csv',\n",
         "      text_column='sentence',\n",
         "      label_column='label',\n",
         "      model_spec=spec,\n",
-        "      delimiter='\\t',\n",
         "      is_training=False)"
       ]
     },
@@ -257,7 +290,9 @@
         "id": "2uZkLR6N6gDR"
       },
       "source": [
-        "Step 3. Customize the TensorFlow model."
+        "**Step 3. Train the TensorFlow model with the training data.**\n",
+        "\n",
+        "The average word embedding model use `batch_size = 32` by default. Therefore you will see that it takes 2104 steps to go through the 67,349 sentences in the training dataset. We will train the model for 10 epochs, which means going through the training dataset 10 times."
       ]
     },
     {
@@ -268,7 +303,7 @@
       },
       "outputs": [],
       "source": [
-        "model = text_classifier.create(train_data, model_spec=spec)"
+        "model = text_classifier.create(train_data, model_spec=spec, epochs=10)"
       ]
     },
     {
@@ -277,7 +312,11 @@
         "id": "-BzCHLWJ6h7q"
       },
       "source": [
-        "Step 4. Evaluate the model."
+        "**Step 4. Evaluate the model with the test data.**\n",
+        "\n",
+        "After training the text classification model using the sentences in the training dataset, we will use the remaining 872 sentences in the test dataset to evaluate how the model perform against new data it has never seen before.\n",
+        "\n",
+        "As the default batch size is 32, it will take 28 steps to go through the 872 sentences in the test dataset."
       ]
     },
     {
@@ -297,21 +336,11 @@
         "id": "CgCDMe0e6jlT"
       },
       "source": [
-        "Step 5.  Export as a TensorFlow Lite model with [metadata](https://www.tensorflow.org/lite/convert/metadata).\n",
+        "**Step 5.  Export as a TensorFlow Lite model.**\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"
+        "Let's export the text classification that we have trained in the TensorFlow Lite format. We will specify which folder to export the model.\n",
+        "\n",
+        "You may see an warning about `vocab.txt` file does not exist in the metadata but they can be safely ignore."
       ]
     },
     {
@@ -322,7 +351,7 @@
       },
       "outputs": [],
       "source": [
-        "model.export(export_dir='mobilebert/', quantization_config=config)"
+        "model.export(export_dir='average_word_vec')"
       ]
     },
     {
@@ -331,9 +360,17 @@
         "id": "rVxaf3x_7OfB"
       },
       "source": [
-        "You can also download the model using the left sidebar in Colab.\n",
+        "You can download the TensorFlow Lite model file using the left sidebar of Colab. Go into the `average_word_vec` folder as we specified in `export_dir` parameter above, right-click on the `model.tflite` file and choose `Download` to download it to your local computer.\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)."
+        "This model can be integrated into an Android or an iOS app using the [NLClassifier API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/nl_classifier) of the [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview).\n",
+        "\n",
+        "See the [TFLite Text Classification sample app](https://github.com/tensorflow/examples/blob/master/lite/examples/text_classification/android/lib_task_api/src/main/java/org/tensorflow/lite/examples/textclassification/client/TextClassificationClient.java#L54) for more details on how the model is used in an working app.\n",
+        "\n",
+        "*Note 1: Android Studio Model Binding does not support text classification yet so please use the TensorFlow Lite Task Library.*\n",
+        "\n",
+        "*Note 2: There is a `model.json` file in the same folder with the TFLite model. It contains the JSON representation of the [metadata](https://www.tensorflow.org/lite/convert/metadata) bundled inside the TensorFlow Lite model. Model metadata helps the TFLite Task Library know what the model does and how to pre-process/post-process data for the model. You don't need to download the `model.json` file as it is only for informational purpose and its content is already inside the TFLite file.*\n",
+        "\n",
+        "*Note 3: If you train a text classification model using MobileBERT or BERT-Base architecture, you will need to use [BertNLClassifier API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/bert_nl_classifier) instead to integrate the trained model into a mobile app.*"
       ]
     },
     {
@@ -342,7 +379,7 @@
         "id": "l65ctmtW7_FF"
       },
       "source": [
-        "The following sections walk through the example step by step to show more detail."
+        "The following sections walk through the example step by step to show more details."
       ]
     },
     {
@@ -351,17 +388,17 @@
         "id": "kJ_B8fMDOhMR"
       },
       "source": [
-        "## Choose a `model_spec` that Represents a Model for Text Classifier\n",
+        "## Choose a model architecture 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",
+        "| Supported Model          | Name of model_spec      | Model Description                                                                                                     | Model size                                  |\n",
+        "|--------------------------|-------------------------|-----------------------------------------------------------------------------------------------------------------------|---------------------------------------------|\n",
+        "| Averaging Word Embedding | 'average_word_vec'      | Averaging text word embeddings with RELU activation.                                                                  |           \u003c1MB                             |\n",
+        "| MobileBERT               | 'mobilebert_classifier' | 4.3x smaller and 5.5x faster than BERT-Base while achieving competitive results, suitable for on-device applications. | 25MB w/ quantization \u003cbr/\u003e 100MB w/o quantization                                        |\n",
+        "| BERT-Base                | 'bert_classifier'       | Standard BERT model that is widely used in NLP tasks.                                                                 | 300MB |\n",
         "\n",
-        "This tutorial uses a smaller model, `average_word_vec` that you can retrain multiple times to demonstrate the process."
+        "In the quick start, we have used the average word embedding model. Let's switch to [MobileBERT](https://arxiv.org/pdf/2004.02984.pdf) to train a model with higher accuracy."
       ]
     },
     {
@@ -372,7 +409,7 @@
       },
       "outputs": [],
       "source": [
-        "spec = model_spec.get('average_word_vec')"
+        "mb_spec = model_spec.get('mobilebert_classifier')"
       ]
     },
     {
@@ -381,44 +418,24 @@
         "id": "ygEncJxtl-nQ"
       },
       "source": [
-        "## Load Input Data Specific to an On-device ML App\n",
+        "## Load training data\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",
+        "You can upload your own dataset to work through this tutorial. Upload your dataset by using the left sidebar in Colab.\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&token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8',\n",
-        "      extract=True)\n",
-        "data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')"
+        "\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",
+        "\n",
+        "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": "E051HBUM5owi"
+        "id": "mWAusqz-WD5i"
       },
       "source": [
-        "The SST-2 dataset has `train.tsv` for training and `dev.tsv` for validation. The files have the following format:\n",
+        "To keep it simple, we will reuse the SST-2 dataset downloaded earlier. Let's use the `TestClassifierDataLoader.from_csv` method to load the data.\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."
+        "Please be noted that as we have changed the model architecture, we will need to reload the training and test dataset to apply the new preprocessing logic."
       ]
     },
     {
@@ -430,18 +447,16 @@
       "outputs": [],
       "source": [
         "train_data = TextClassifierDataLoader.from_csv(\n",
-        "      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),\n",
+        "      filename='train.csv',\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",
+        "      filename='dev.csv',\n",
         "      text_column='sentence',\n",
         "      label_column='label',\n",
         "      model_spec=spec,\n",
-        "      delimiter='\\t',\n",
         "      is_training=False)"
       ]
     },
@@ -460,9 +475,11 @@
         "id": "AWuoensX4vDA"
       },
       "source": [
-        "## Customize the TensorFlow Model\n",
+        "## Train a TensorFlow Model\n",
         "\n",
-        "Create a custom text classifier model based on the loaded data."
+        "Train a text classification model using the training data.\n",
+        "\n",
+        "*Note: As MobileBERT is a complex model, each training epoch will takes about 10 minutes on a Colab GPU. Please make sure that you are using a GPU runtime.* "
       ]
     },
     {
@@ -473,7 +490,7 @@
       },
       "outputs": [],
       "source": [
-        "model = text_classifier.create(train_data, model_spec=spec, epochs=10)"
+        "model = text_classifier.create(train_data, model_spec=spec, epochs=3)"
       ]
     },
     {
@@ -502,9 +519,9 @@
         "id": "LP5FPk_tOxoZ"
       },
       "source": [
-        "## Evaluate the Customized Model\n",
+        "## Evaluate the model\n",
         "\n",
-        "Evaluate the model with the test data and get its loss and accuracy."
+        "Evaluate the model that we have just trained using the test data and measure the loss and accuracy value."
       ]
     },
     {
@@ -524,9 +541,32 @@
         "id": "aeHoGAceO2xV"
       },
       "source": [
-        "## Export as a TensorFlow Lite Model\n",
+        "## Quantize the 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`."
+        "In many on-device ML application, the model size is an important factor. Therefore, it is recommended that you apply quantize the model to make it smaller and potentially run faster. Model Maker automatically applies the recommended quantization scheme for each model architecture but you can customize the quantization config as below."
+      ]
+    },
+    {
+      "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": "markdown",
+      "metadata": {
+        "id": "esBGwHE2QxE8"
+      },
+      "source": [
+        "## Export as a TensorFlow Lite model\n",
+        "\n",
+        "Convert the trained model to TensorFlow Lite model format with [metadata](https://www.tensorflow.org/lite/convert/metadata) so 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`."
       ]
     },
     {
@@ -537,7 +577,7 @@
       },
       "outputs": [],
       "source": [
-        "model.export(export_dir='average_word_vec/')"
+        "model.export(export_dir='mobilebert/', quantization_config=config)"
       ]
     },
     {
@@ -546,7 +586,7 @@
         "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)."
+        "The TensorFlow Lite model file can be integrated in a mobile app using the [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). Please note that this is **different** from the `NLClassifier` API used to integrate the text classification trained with the average word vector model architecture."
       ]
     },
     {
@@ -555,14 +595,14 @@
         "id": "AVy0ormoMZwL"
       },
       "source": [
-        "The allowed export formats can be one or a list of the following:\n",
+        "The 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:"
+        "By default, it exports only the TensorFlow Lite model file containing the model metadata. You can also choose to export other files related to the model for better examination. For instance, exporting only the label file and vocab file as follows:"
       ]
     },
     {
@@ -573,7 +613,7 @@
       },
       "outputs": [],
       "source": [
-        "model.export(export_dir='average_word_vec/', export_format=[ExportFormat.LABEL, ExportFormat.VOCAB])"
+        "model.export(export_dir='mobilebert/', export_format=[ExportFormat.LABEL, ExportFormat.VOCAB])"
       ]
     },
     {
@@ -582,7 +622,7 @@
         "id": "HZKYthlVrTos"
       },
       "source": [
-        "You can evaluate the tflite model with `evaluate_tflite` method to get its accuracy."
+        "You can evaluate the TFLite model with `evaluate_tflite` method to measure its accuracy. Converting the trained TensorFlow model to TFLite format and apply quantization can affect its accuracy so it is recommended to evaluate the TFLite model accuracy before deployment."
       ]
     },
     {
@@ -593,7 +633,8 @@
       },
       "outputs": [],
       "source": [
-        "accuracy = model.evaluate_tflite('average_word_vec/model.tflite', test_data)"
+        "accuracy = model.evaluate_tflite('mobilebert/model.tflite', test_data)\n",
+        "print('TFLite model accuracy: ', accuracy)"
       ]
     },
     {
@@ -615,10 +656,46 @@
     {
       "cell_type": "markdown",
       "metadata": {
+        "id": "E8VxPiOLy4Gv"
+      },
+      "source": [
+        "### Customize the MobileBERT model hyperparameters\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": "mwtiksguDfhl"
       },
       "source": [
-        "### Adjust the model\n",
+        "### Customize the average word embedding model hyperparameters\n",
         "\n",
         "You can adjust the model infrastructure like the `wordvec_dim` and the `seq_len` variables in the `AverageWordVecModelSpec` class.\n"
       ]
@@ -661,11 +738,10 @@
       "outputs": [],
       "source": [
         "new_train_data = TextClassifierDataLoader.from_csv(\n",
-        "      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),\n",
+        "      filename='train.csv',\n",
         "      text_column='sentence',\n",
         "      label_column='label',\n",
         "      model_spec=new_model_spec,\n",
-        "      delimiter='\\t',\n",
         "      is_training=True)"
       ]
     },
@@ -692,42 +768,6 @@
     {
       "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": [
@@ -748,7 +788,7 @@
       },
       "outputs": [],
       "source": [
-        "model = text_classifier.create(train_data, model_spec=spec, epochs=20)"
+        "model = text_classifier.create(new_train_data, model_spec=new_model_spec, epochs=20)"
       ]
     },
     {
@@ -768,7 +808,14 @@
       },
       "outputs": [],
       "source": [
-        "loss, accuracy = model.evaluate(test_data)"
+        "new_test_data = TextClassifierDataLoader.from_csv(\n",
+        "      filename='dev.csv',\n",
+        "      text_column='sentence',\n",
+        "      label_column='label',\n",
+        "      model_spec=new_model_spec,\n",
+        "      is_training=False)\n",
+        "\n",
+        "loss, accuracy = model.evaluate(new_test_data)"
       ]
     },
     {
@@ -809,7 +856,17 @@
     "accelerator": "GPU",
     "colab": {
       "collapsed_sections": [],
-      "name": "model_maker_text_classification.ipynb",
+      "name": "Model Maker Text Classification Tutorial",
+      "provenance": [
+        {
+          "file_id": "1dbRXQCjtm-jBFC32DJ6YCVXnXBOG3M5t",
+          "timestamp": 1613441434239
+        },
+        {
+          "file_id": "https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/tutorials/model_maker_text_classification.ipynb",
+          "timestamp": 1612303859066
+        }
+      ],
       "toc_visible": true
     },
     "kernelspec": {