#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2019 The TensorFlow Authors.
End of explanation
"""


!pip install tflite-model-maker

"""
Explanation: Text classification with TensorFlow Lite Model Maker
<table class="tfo-notebook-buttons" align="left">
  <td>
    <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>
  </td>
  <td>
    <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>
  </td>
  <td>
    <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>
  </td>
  <td>
    <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>
  </td>
</table>

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.
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.
Prerequisites
Install the required packages
To run this example, install the required packages, including the Model Maker package from the GitHub repo.
End of explanation
"""


import numpy as np
import os

import tensorflow as tf
assert tf.__version__.startswith('2')

from tflite_model_maker import configs
from tflite_model_maker import ExportFormat
from tflite_model_maker import model_spec
from tflite_model_maker import text_classifier
from tflite_model_maker import TextClassifierDataLoader

"""
Explanation: Import the required packages.
End of explanation
"""


data_dir = tf.keras.utils.get_file(
      fname='SST-2.zip',
      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',
      extract=True)
data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')

"""
Explanation: Get the data path
Download the dataset for this tutorial.
End of explanation
"""


spec = model_spec.get('mobilebert_classifier')

"""
Explanation: You can also upload your own dataset to work through this tutorial. Upload your dataset by using the left sidebar in Colab.
<img src="https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_text_classification.png" alt="Upload File" width="800" hspace="100">
If you prefer not to upload your dataset to the cloud, you can also locally run the library by following the guide.
End-to-End Workflow
This workflow consists of five steps as outlined below:
Step 1. Choose a model specification that represents a text classification model.
This tutorial uses MobileBERT as an example.
End of explanation
"""


train_data = TextClassifierDataLoader.from_csv(
      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      delimiter='\t',
      is_training=True)
test_data = TextClassifierDataLoader.from_csv(
      filename=os.path.join(os.path.join(data_dir, 'dev.tsv')),
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      delimiter='\t',
      is_training=False)

"""
Explanation: Step 2.   Load train and test data specific to an on-device ML app and preprocess the data according to a specific model_spec.
End of explanation
"""


model = text_classifier.create(train_data, model_spec=spec)

"""
Explanation: Step 3. Customize the TensorFlow model.
End of explanation
"""


loss, acc = model.evaluate(test_data)

"""
Explanation: Step 4. Evaluate the model.
End of explanation
"""


config = configs.QuantizationConfig.create_dynamic_range_quantization(optimizations=[tf.lite.Optimize.OPTIMIZE_FOR_LATENCY])
config._experimental_new_quantizer = True

model.export(export_dir='mobilebert/', quantization_config=config)

"""
Explanation: Step 5.  Export as a TensorFlow Lite model with metadata.
Since MobileBERT is too big for on-device applications, use dynamic range quantization on the model to compress it by almost 4x with minimal performance degradation.
End of explanation
"""


spec = model_spec.get('average_word_vec')

"""
Explanation: You can also download the model using the left sidebar in Colab.
After executing the 5 steps above, you can further use the TensorFlow Lite model file in on-device applications using BertNLClassifier API in TensorFlow Lite Task Library.
The following sections walk through the example step by step to show more detail.
Choose a model_spec that Represents a Model for Text Classifier
Each model_spec object represents a specific model for the text classifier. TensorFlow Lite Model Maker currently supports MobileBERT, averaging word embeddings and [BERT-Base]((https://arxiv.org/pdf/1810.04805.pdf) models.
Supported Model | Name of model_spec | Model Description
--- | --- | ---
MobileBERT | 'mobilebert_classifier' | 4.3x smaller and 5.5x faster than BERT-Base while achieving competitive results, suitable for on-device applications.
BERT-Base | 'bert_classifier' | Standard BERT model that is widely used in NLP tasks.
averaging word embedding | 'average_word_vec' | Averaging text word embeddings with RELU activation.
This tutorial uses a smaller model, average_word_vec that you can retrain multiple times to demonstrate the process.
End of explanation
"""


data_dir = tf.keras.utils.get_file(
      fname='SST-2.zip',
      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',
      extract=True)
data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')

"""
Explanation: Load Input Data Specific to an On-device ML App
The SST-2 (Stanford Sentiment Treebank) is one of the tasks in the GLUE 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.
Download the archived version of the dataset and extract it.
End of explanation
"""


train_data = TextClassifierDataLoader.from_csv(
      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      delimiter='\t',
      is_training=True)
test_data = TextClassifierDataLoader.from_csv(
      filename=os.path.join(os.path.join(data_dir, 'dev.tsv')),
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      delimiter='\t',
      is_training=False)

"""
Explanation: The SST-2 dataset has train.tsv for training and dev.tsv for validation. The files have the following format:
sentence | label
--- | ---
it 's a charming and often affecting journey . | 1
unflinchingly bleak and desperate | 0
A positive review is labeled 1 and a negative review is labeled 0.
Use the TestClassifierDataLoader.from_csv method to load the data.
End of explanation
"""


model = text_classifier.create(train_data, model_spec=spec, epochs=10)

"""
Explanation: 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.
Customize the TensorFlow Model
Create a custom text classifier model based on the loaded data.
End of explanation
"""


model.summary()

"""
Explanation: Examine the detailed model structure.
End of explanation
"""


loss, acc = model.evaluate(test_data)

"""
Explanation: Evaluate the Customized Model
Evaluate the result of the model and get the loss and accuracy of the model.
Evaluate the loss and accuracy in the test data.
End of explanation
"""


model.export(export_dir='average_word_vec/')

"""
Explanation: Export as a TensorFlow Lite Model
Convert the existing model to TensorFlow Lite model format with 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.
End of explanation
"""


model.export(export_dir='average_word_vec/', export_format=[ExportFormat.LABEL, ExportFormat.VOCAB])

"""
Explanation: The TensorFlow Lite model file can be used in the text classification reference app using NLClassifier API in TensorFlow Lite Task Library.
The allowed export formats can be one or a list of the following:

ExportFormat.TFLITE
ExportFormat.LABEL
ExportFormat.VOCAB
ExportFormat.SAVED_MODEL

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:
End of explanation
"""


model.evaluate_tflite('average_word_vec/model.tflite', test_data)

"""
Explanation: You can evalute the tflite model with evaluate_tflite method.
End of explanation
"""


new_model_spec = model_spec.AverageWordVecModelSpec(wordvec_dim=32)

"""
Explanation: Advanced Usage
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:

Creates the model for the text classifier according to model_spec.
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.

This section covers advanced usage topics like adjusting the model and the training hyperparameters.
Adjust the model
You can adjust the model infrastructure like the wordvec_dim and the seq_len variables in the AverageWordVecModelSpec class.
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.
End of explanation
"""


new_train_data = TextClassifierDataLoader.from_csv(
      filename=os.path.join(os.path.join(data_dir, 'train.tsv')),
      text_column='sentence',
      label_column='label',
      model_spec=new_model_spec,
      delimiter='\t',
      is_training=True)

"""
Explanation: Get the preprocessed data.
End of explanation
"""


model = text_classifier.create(new_train_data, model_spec=new_model_spec)

"""
Explanation: Train the new model.
End of explanation
"""


new_model_spec = model_spec.get('mobilebert_classifier')
new_model_spec.seq_len = 256

"""
Explanation: You can also adjust the MobileBERT model.
The model parameters you can adjust are:

seq_len: Length of the sequence to feed into the model.
initializer_range: The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
trainable: Boolean that specifies whether the pre-trained layer is trainable.

The training pipeline parameters you can adjust are:

model_dir: The location of the model checkpoint files. If not set, a temporary directory will be used.
dropout_rate: The dropout rate.
learning_rate: The initial learning rate for the Adam optimizer.
tpu: TPU address to connect to.

For instance, you can set the seq_len=256 (default is 128). This allows the model to classify longer text.
End of explanation
"""


model = text_classifier.create(train_data, model_spec=spec, epochs=20)

"""
Explanation: Tune the training hyperparameters
You can also tune the training hyperparameters like epochs and batch_size that affect the model accuracy. For instance,

epochs: more epochs could achieve better accuracy, but may lead to overfitting.
batch_size: the number of samples to use in one training step.

For example, you can train with more epochs.
End of explanation
"""


loss, accuracy = model.evaluate(test_data)

"""
Explanation: Evaluate the newly retrained model with 20 training epochs.
End of explanation
"""


spec = model_spec.get('bert_classifier')

"""
Explanation: Change the Model Architecture
You can change the model by changing the model_spec. The following shows how to change to BERT-Base model.
Change the model_spec to BERT-Base model for the text classifier.
End of explanation
"""


