classify_text_with_bert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IZ6SNYq_tVVC"
   },
   "source": [
    "# Classify text with BERT\n",
    "\n",
    "### Learning Objectives\n",
    "1. Learn how to load a pre-trained BERT model from TensorFlow Hub\n",
    "2. Learn how to build your own model by combining with a classifier\n",
    "3. Learn how to train a BERT model by fine-tuning\n",
    "4. Learn how to save your trained model and use it\n",
    "5. Learn how to evaluate a text classification model\n",
    "\n",
    "This lab will show you how to fine-tune BERT to perform sentiment analysis on a dataset of plain-text IMDB movie reviews.\n",
    "In addition to training a model, you will learn how to preprocess text into an appropriate format.\n",
    "\n",
    "### Before you start\n",
    "\n",
    "Please ensure you have a GPU (1 x NVIDIA Tesla T4 should be enough) attached to your Notebook instance to ensure that the training doesn't take too long. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2PHBpLPuQdmK"
   },
   "source": [
    "## About BERT\n",
    "\n",
    "[BERT](https://arxiv.org/abs/1810.04805) and other Transformer encoder architectures have been wildly successful on a variety of tasks in NLP (natural language processing). They compute vector-space representations of natural language that are suitable for use in deep learning models. The BERT family of models uses the Transformer encoder architecture to process each token of input text in the full context of all tokens before and after, hence the name: Bidirectional Encoder Representations from Transformers. \n",
    "\n",
    "BERT models are usually pre-trained on a large corpus of text, then fine-tuned for specific tasks.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SCjmX4zTCkRK"
   },
   "source": [
    "### Setup\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select the `bert_kernel` for this Jupyter notebook from the kernel menu in the top-right corner.\n",
    "\n",
    "If the `bert_kernel` is unavailable, run the following cell to create it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ~/asl-ml-immersion && make bert_kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_XgTpm9ZxoN9",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow_text as text\n",
    "from google.cloud import aiplatform\n",
    "from official.nlp import optimization  # to create AdamW optmizer\n",
    "\n",
    "tf.get_logger().setLevel(\"ERROR\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check if you have a GPU attached. Run the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Num GPUs Available: \", len(tf.config.list_physical_devices(\"GPU\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q6MugfEgDRpY"
   },
   "source": [
    "### Sentiment Analysis\n",
    "\n",
    "This notebook trains a sentiment analysis model to classify movie reviews as *positive* or *negative*, based on the text of the review.\n",
    "\n",
    "You'll use the [Large Movie Review Dataset](https://ai.stanford.edu/~amaas/data/sentiment/) that contains the text of 50,000 movie reviews from the [Internet Movie Database](https://www.imdb.com/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vnvd4mrtPHHV"
   },
   "source": [
    "### Download the IMDB dataset\n",
    "\n",
    "Let's download and extract the dataset, then explore the directory structure.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pOdqCMoQDRJL",
    "tags": []
   },
   "outputs": [],
   "source": [
    "url = \"https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\"\n",
    "\n",
    "# Set a path to a folder outside the git repo. This is important so data won't get indexed by git on Jupyter lab\n",
    "path = \"/home/jupyter/\"\n",
    "\n",
    "dataset = tf.keras.utils.get_file(\n",
    "    \"aclImdb_v1.tar.gz\", url, untar=True, cache_dir=path, cache_subdir=\"\"\n",
    ")\n",
    "\n",
    "dataset_dir = os.path.join(os.path.dirname(dataset), \"aclImdb\")\n",
    "\n",
    "train_dir = os.path.join(dataset_dir, \"train\")\n",
    "\n",
    "# remove unused folders to make it easier to load the data\n",
    "remove_dir = os.path.join(train_dir, \"unsup\")\n",
    "shutil.rmtree(remove_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lN9lWCYfPo7b"
   },
   "source": [
    "Next, you will use the `text_dataset_from_directory` utility to create a labeled `tf.data.Dataset`.\n",
    "\n",
    "The IMDB dataset has already been divided into train and test, but it lacks a validation set. Let's create a validation set using an 80:20 split of the training data by using the `validation_split` argument below.\n",
    "\n",
    "Note:  When using the `validation_split` and `subset` arguments, make sure to either specify a random seed, or to pass `shuffle=False`, so that the validation and training splits have no overlap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6IwI_2bcIeX8",
    "tags": []
   },
   "outputs": [],
   "source": [
    "AUTOTUNE = tf.data.AUTOTUNE\n",
    "batch_size = 32\n",
    "seed = 42\n",
    "\n",
    "raw_train_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "    path + \"aclImdb/train\",\n",
    "    batch_size=batch_size,\n",
    "    validation_split=0.2,\n",
    "    subset=\"training\",\n",
    "    seed=seed,\n",
    ")\n",
    "\n",
    "class_names = raw_train_ds.class_names\n",
    "train_ds = raw_train_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "\n",
    "val_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "    path + \"aclImdb/train\",\n",
    "    batch_size=batch_size,\n",
    "    validation_split=0.2,\n",
    "    subset=\"validation\",\n",
    "    seed=seed,\n",
    ")\n",
    "\n",
    "val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "\n",
    "test_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "    path + \"aclImdb/test\", batch_size=batch_size\n",
    ")\n",
    "\n",
    "test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HGm10A5HRGXp"
   },
   "source": [
    "Let's take a look at a few reviews."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JuxDkcvVIoev",
    "tags": []
   },
   "outputs": [],
   "source": [
    "for text_batch, label_batch in train_ds.take(1):\n",
    "    for i in range(3):\n",
    "        print(f\"Review: {text_batch.numpy()[i]}\")\n",
    "        label = label_batch.numpy()[i]\n",
    "        print(f\"Label : {label} ({class_names[label]})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dX8FtlpGJRE6"
   },
   "source": [
    "## Loading models from TensorFlow Hub\n",
    "\n",
    "For the purpose of this lab, we will be loading a model called Small BERT. Small BERT has the same general architecture as the original BERT but the has fewer and/or smaller Transformer blocks. \n",
    "\n",
    "Some other popular BERT models are BERT Base, ALBERT, BERT Experts, Electra. See the continued learning section at the end of this lab for more info. \n",
    "\n",
    "Aside from the models available below, there are [multiple versions](https://tfhub.dev/google/collections/transformer_encoders_text/1) of the models that are larger and can yeld even better accuracy but they are too big to be fine-tuned on a single GPU. You will be able to do that on the [Solve GLUE tasks using BERT on a TPU colab](https://www.tensorflow.org/tutorials/text/solve_glue_tasks_using_bert_on_tpu).\n",
    "\n",
    "You'll see in the code below that switching the tfhub.dev URL is enough to try any of these models, because all the differences between them are encapsulated in the SavedModels from TF Hub."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose a BERT model to fine-tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "y8_ctG55-uTX",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# defining the URL of the smallBERT model to use\n",
    "tfhub_handle_encoder = (\n",
    "    \"https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/1\"\n",
    ")\n",
    "\n",
    "# defining the corresponding preprocessing model for the BERT model above\n",
    "tfhub_handle_preprocess = (\n",
    "    \"https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3\"\n",
    ")\n",
    "\n",
    "print(f\"BERT model selected           : {tfhub_handle_encoder}\")\n",
    "print(f\"Preprocess model auto-selected: {tfhub_handle_preprocess}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7WrcxxTRDdHi"
   },
   "source": [
    "## Preprocessing model\n",
    "\n",
    "Text inputs need to be transformed to numeric token ids and arranged in several Tensors before being input to BERT. TensorFlow Hub provides a matching preprocessing model for each of the BERT models discussed above, which implements this transformation using TF ops from the TF.text library. It is not necessary to run pure Python code outside your TensorFlow model to preprocess text.\n",
    "\n",
    "The preprocessing model must be the one referenced by the documentation of the BERT model, which you can read at the URL printed above. For BERT models from the drop-down above, the preprocessing model is selected automatically.\n",
    "\n",
    "Note: You will load the preprocessing model into a [hub.KerasLayer](https://www.tensorflow.org/hub/api_docs/python/hub/KerasLayer) to compose your fine-tuned model. This is the preferred API to load a TF2-style SavedModel from TF Hub into a Keras model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `hub.KerasLayer` to initialize the preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_preprocess_model = hub.KerasLayer(tfhub_handle_preprocess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x4naBiEE_cZX"
   },
   "source": [
    "Let's try the preprocessing model on some text and see the output:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the preprocess model function and pass text_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "r9-zCzJpnuwS",
    "tags": []
   },
   "outputs": [],
   "source": [
    "text_test = [\"this is such an amazing movie!\"]\n",
    "text_preprocessed = bert_preprocess_model(text_test)\n",
    "\n",
    "# This print box will help you inspect the keys in the pre-processed dictionary\n",
    "print(f\"Keys       : {list(text_preprocessed.keys())}\")\n",
    "\n",
    "# 1. input_word_ids is the ids for the words in the tokenized sentence\n",
    "print(f'Shape      : {text_preprocessed[\"input_word_ids\"].shape}')\n",
    "print(f'Word Ids   : {text_preprocessed[\"input_word_ids\"][0, :12]}')\n",
    "\n",
    "# 2. input_mask is the tokens which we are masking (masked language model)\n",
    "print(f'Input Mask : {text_preprocessed[\"input_mask\"][0, :12]}')\n",
    "\n",
    "# 3. input_type_ids is the sentence id of the input sentence.\n",
    "print(f'Type Ids   : {text_preprocessed[\"input_type_ids\"][0, :12]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EqL7ihkN_862"
   },
   "source": [
    "As you can see, now you have the 3 outputs from the preprocessing that a BERT model would use (`input_words_id`, `input_mask` and `input_type_ids`).\n",
    "\n",
    "Some other important points:\n",
    "- The input is truncated to 128 tokens. \n",
    "- The `input_type_ids` only have one value (0) because this is a single sentence input. For a multiple sentence input, it would have one number for each input.\n",
    "\n",
    "The text pre-processor is a TensorFlow model. This means that instead of pre-processing separately, we can include it as a layer in the model code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DKnLPSEmtp9i"
   },
   "source": [
    "### Using the BERT model\n",
    "\n",
    "Before putting BERT into your own model, let's take a look at its outputs. You will load it from TF Hub and see the returned values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tXxYpK8ixL34",
    "tags": []
   },
   "outputs": [],
   "source": [
    "bert_model = hub.KerasLayer(tfhub_handle_encoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_OoF9mebuSZc",
    "tags": []
   },
   "outputs": [],
   "source": [
    "bert_results = bert_model(text_preprocessed)\n",
    "\n",
    "print(f\"Loaded BERT: {tfhub_handle_encoder}\")\n",
    "print(f'Pooled Outputs Shape:{bert_results[\"pooled_output\"].shape}')\n",
    "print(f'Pooled Outputs Values:{bert_results[\"pooled_output\"][0, :12]}')\n",
    "print(f'Sequence Outputs Shape:{bert_results[\"sequence_output\"].shape}')\n",
    "print(f'Sequence Outputs Values:{bert_results[\"sequence_output\"][0, :12]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sm61jDrezAll"
   },
   "source": [
    "The BERT models return a map with 3 important keys: `pooled_output`, `sequence_output`, `encoder_outputs`:\n",
    "\n",
    "- `pooled_output` to represent each input sequence as a whole. The shape is `[batch_size, H]`. You can think of this as an embedding for the entire movie review.\n",
    "- `sequence_output` represents each input token in the context. The shape is `[batch_size, seq_length, H]`. You can think of this as a contextual embedding for every token in the movie review.\n",
    "- `encoder_outputs` are the intermediate activations of the `L` Transformer blocks. `outputs[\"encoder_outputs\"][i]` is a Tensor of shape `[batch_size, seq_length, 1024]` with the outputs of the i-th Transformer block, for `0 <= i < L`. The last value of the list is equal to `sequence_output`.\n",
    "\n",
    "For the fine-tuning you are going to use the `pooled_output` array."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pDNKfAXbDnJH"
   },
   "source": [
    "### Define your model\n",
    "\n",
    "You will create a very simple fine-tuned model, with the preprocessing model, the selected BERT model, one Dense and a Dropout layer.\n",
    "\n",
    "Note: for more information about the base model's input and output you can use copy the model's url to get to the documentation page."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The order of the layers in the model will be:\n",
    "1. Input Layer\n",
    "2. Pre-processing Layer\n",
    "3. Encoder Layer\n",
    "4. From the BERT output map, use pooled_output\n",
    "5. Dropout layer\n",
    "6. Dense layer with sigmoid activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aksj743St9ga",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def build_classifier_model(dropout_rate=0.1):\n",
    "    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name=\"text\")\n",
    "    preprocessing_layer = hub.KerasLayer(\n",
    "        tfhub_handle_preprocess, name=\"preprocessing\"\n",
    "    )\n",
    "    encoder_inputs = preprocessing_layer(text_input)\n",
    "    encoder = hub.KerasLayer(\n",
    "        tfhub_handle_encoder, trainable=True, name=\"BERT_encoder\"\n",
    "    )\n",
    "    outputs = encoder(encoder_inputs)\n",
    "    net = outputs[\"pooled_output\"]\n",
    "    net = tf.keras.layers.Dropout(dropout_rate)(net)\n",
    "    net = tf.keras.layers.Dense(1, activation=\"sigmoid\", name=\"classifier\")(net)\n",
    "    return tf.keras.Model(text_input, net)\n",
    "\n",
    "\n",
    "# Let's check that the model runs with the output of the preprocessing model.\n",
    "dropout_rate = 0.15\n",
    "classifier_model = build_classifier_model(dropout_rate)\n",
    "bert_raw_result = classifier_model(tf.constant(text_test))\n",
    "print(bert_raw_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZTUzNV2JE2G3"
   },
   "source": [
    "The output is meaningless, of course, because the model has not been trained yet.\n",
    "\n",
    "Let's take a look at the model's structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0EmzyHZXKIpm",
    "tags": []
   },
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(classifier_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WbUWoZMwc302"
   },
   "source": [
    "## Model training\n",
    "\n",
    "You now have all the pieces to train a model, including the preprocessing module, BERT encoder, data, and classifier."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WpJ3xcwDT56v"
   },
   "source": [
    "### Loss function\n",
    "\n",
    "Since this is a binary classification problem and the model outputs a probability (a single-unit layer), you'll use `losses.BinaryCrossentropy` loss function.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define your loss and evaluation metrics here. Since it is a binary classification use BinaryCrossentropy and BinaryAccuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OWPOZE-L3AgE",
    "tags": []
   },
   "outputs": [],
   "source": [
    "loss = tf.keras.losses.BinaryCrossentropy()\n",
    "metrics = tf.metrics.BinaryAccuracy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "77psrpfzbxtp"
   },
   "source": [
    "### Optimizer\n",
    "\n",
    "For fine-tuning, let's use the same optimizer that BERT was originally trained with: the \"Adaptive Moments\" (Adam). This optimizer minimizes the prediction loss and does regularization by weight decay (not using moments), which is also known as [AdamW](https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/AdamW).\n",
    "\n",
    "In past labs, we have been using the Adam optimizer which is a popular choice. However, for this lab we will be using a new optimizier which is meant to improve generalization. The intuition and algoritm behind AdamW can be found in this paper [here](https://arxiv.org/abs/1711.05101).\n",
    "\n",
    "For the learning rate (`init_lr`), we use the same schedule as BERT pre-training: linear decay of a notional initial learning rate, prefixed with a linear warm-up phase over the first 10% of training steps (`num_warmup_steps`). In line with the BERT paper, the initial learning rate is smaller for fine-tuning (best of 5e-5, 3e-5, 2e-5)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P9eP2y9dbw32",
    "tags": []
   },
   "outputs": [],
   "source": [
    "epochs = 5\n",
    "steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()\n",
    "num_train_steps = steps_per_epoch * epochs\n",
    "num_warmup_steps = int(0.1 * num_train_steps)\n",
    "\n",
    "init_lr = 3e-5\n",
    "optimizer = optimization.create_optimizer(\n",
    "    init_lr=init_lr,\n",
    "    num_train_steps=num_train_steps,\n",
    "    num_warmup_steps=num_warmup_steps,\n",
    "    optimizer_type=\"adamw\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SqlarlpC_v0g"
   },
   "source": [
    "### Loading the BERT model and training\n",
    "\n",
    "Using the `classifier_model` you created earlier, you can compile the model with the loss, metric and optimizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Complile the model using the optimizer, loss and metrics you defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-7GPDhR98jsD",
    "tags": []
   },
   "outputs": [],
   "source": [
    "classifier_model.compile(optimizer=optimizer, loss=loss, metrics=metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CpBuV5j2cS_b"
   },
   "source": [
    "Note: training time will vary depending on the complexity of the BERT model you have selected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's train the model for a few epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HtfDFAnN_Neu"
   },
   "outputs": [],
   "source": [
    "print(f\"Training model with {tfhub_handle_encoder}\")\n",
    "history = classifier_model.fit(\n",
    "    x=train_ds, validation_data=val_ds, epochs=epochs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uBthMlTSV8kn"
   },
   "source": [
    "### Evaluate the model\n",
    "\n",
    "Let's see how the model performs. Two values will be returned. Loss (a number which represents the error, lower values are better), and accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "slqB-urBV9sP"
   },
   "outputs": [],
   "source": [
    "loss, accuracy = classifier_model.evaluate(test_ds)\n",
    "\n",
    "print(f\"Loss: {loss}\")\n",
    "print(f\"Accuracy: {accuracy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uttWpgmSfzq9"
   },
   "source": [
    "### Plot the accuracy and loss over time\n",
    "\n",
    "Based on the `History` object returned by `model.fit()`. You can plot the training and validation loss for comparison, as well as the training and validation accuracy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fiythcODf0xo"
   },
   "outputs": [],
   "source": [
    "history_dict = history.history\n",
    "print(history_dict.keys())\n",
    "\n",
    "acc = history_dict[\"binary_accuracy\"]\n",
    "val_acc = history_dict[\"val_binary_accuracy\"]\n",
    "loss = history_dict[\"loss\"]\n",
    "val_loss = history_dict[\"val_loss\"]\n",
    "\n",
    "epochs = range(1, len(acc) + 1)\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.subplot(2, 1, 1)\n",
    "# \"bo\" is for \"blue dot\"\n",
    "plt.plot(epochs, loss, \"r\", label=\"Training loss\")\n",
    "# b is for \"solid blue line\"\n",
    "plt.plot(epochs, val_loss, \"b\", label=\"Validation loss\")\n",
    "plt.title(\"Training and validation loss\")\n",
    "# plt.xlabel('Epochs')\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(epochs, acc, \"r\", label=\"Training acc\")\n",
    "plt.plot(epochs, val_acc, \"b\", label=\"Validation acc\")\n",
    "plt.title(\"Training and validation accuracy\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend(loc=\"lower right\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WzJZCo-cf-Jf"
   },
   "source": [
    "In this plot, the red lines represents the training loss and accuracy, and the blue lines are the validation loss and accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Rtn7jewb6dg4"
   },
   "source": [
    "### Export for inference\n",
    "\n",
    "Now you just save your fine-tuned model for later use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the model to saved_model_path for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ShcvqJAgVera"
   },
   "outputs": [],
   "source": [
    "dataset_name = \"imdb\"\n",
    "saved_model_path = \"./{}_bert\".format(dataset_name.replace(\"/\", \"_\"))\n",
    "TIMESTAMP = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
    "\n",
    "EXPORT_PATH = os.path.join(saved_model_path, TIMESTAMP)\n",
    "\n",
    "classifier_model.save(EXPORT_PATH, include_optimizer=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PbI25bS1vD7s"
   },
   "source": [
    "Let's reload the model so you can try it side by side with the model that is still in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gUEWVskZjEF0"
   },
   "outputs": [],
   "source": [
    "reloaded_model = tf.saved_model.load(EXPORT_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oyTappHTvNCz"
   },
   "source": [
    "Here you can test your model on any sentence you want, just add to the examples variable below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VBWzH6exlCPS"
   },
   "outputs": [],
   "source": [
    "def print_my_examples(inputs, results):\n",
    "    result_for_printing = [\n",
    "        f\"input: {inputs[i]:<30} : score: {results[i][0]:.6f}\"\n",
    "        for i in range(len(inputs))\n",
    "    ]\n",
    "    print(*result_for_printing, sep=\"\\n\")\n",
    "    print()\n",
    "\n",
    "\n",
    "examples = [\n",
    "    \"this is such an amazing movie!\",  # this is the same sentence tried earlier\n",
    "    \"The movie was great!\",\n",
    "    \"The movie was meh.\",\n",
    "    \"The movie was okish.\",\n",
    "    \"The movie was terrible...\",\n",
    "]\n",
    "\n",
    "reloaded_results = reloaded_model(tf.constant(examples))\n",
    "original_results = classifier_model(tf.constant(examples))\n",
    "\n",
    "print(\"Results from the saved model:\")\n",
    "print_my_examples(examples, reloaded_results)\n",
    "print(\"Results from the model in memory:\")\n",
    "print_my_examples(examples, original_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## (Optional) Deploy your model on Vertex AI to get online predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3cOmih754Y_M"
   },
   "source": [
    "If you want to use your model on [TF Serving](https://www.tensorflow.org/tfx/guide/serving), remember that it will call your SavedModel through one of its named signatures. In Python, you can test them as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0FdVD3973S-O"
   },
   "outputs": [],
   "source": [
    "serving_results = reloaded_model.signatures[\"serving_default\"](\n",
    "    tf.constant(examples)\n",
    ")\n",
    "\n",
    "serving_results = serving_results[\"classifier\"]\n",
    "\n",
    "print_my_examples(examples, serving_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll export the model to a TensorFlow SavedModel format. Once we have a model in this format, we have lots of ways to \"serve\" the model, from a web application, from JavaScript, from mobile applications, etc.\n",
    "\n",
    "Next, print the signature of your saved model using the SavedModel Command Line Interface command saved_model_cli. You can read more about the command line interface and the show and run commands it supports in the [documentation here](https://www.tensorflow.org/guide/saved_model#overview_of_commands)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show \\\n",
    "    --tag_set serve \\\n",
    "    --signature_def serving_default \\\n",
    "    --dir {EXPORT_PATH}\n",
    "\n",
    "!find {EXPORT_PATH}\n",
    "os.environ['EXPORT_PATH'] = EXPORT_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TIMESTAMP = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
    "PROJECT = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "PROJECT = PROJECT[0]\n",
    "BUCKET = PROJECT\n",
    "REGION = \"us-central1\"\n",
    "MODEL_DISPLAYNAME = f\"classification-bert-{TIMESTAMP}\"\n",
    "\n",
    "print(f\"MODEL_DISPLAYNAME: {MODEL_DISPLAYNAME}\")\n",
    "\n",
    "# from https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers\n",
    "SERVING_CONTAINER_IMAGE_URI = (\n",
    "    \"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest\"\n",
    ")\n",
    "\n",
    "os.environ[\"BUCKET\"] = BUCKET\n",
    "os.environ[\"REGION\"] = REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Create GCS bucket if it doesn't exist already...\n",
    "exists=$(gsutil ls -d | grep -w gs://${BUCKET}/)\n",
    "\n",
    "if [ -n \"$exists\" ]; then\n",
    "    echo -e \"Bucket exists, let's not recreate it.\"\n",
    "else\n",
    "    echo \"Creating a new GCS bucket.\"\n",
    "    gsutil mb -l ${REGION} gs://${BUCKET}\n",
    "    echo \"\\nHere are your current buckets:\"\n",
    "    gsutil ls\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gsutil cp -r $EXPORT_PATH gs://$BUCKET/$MODEL_DISPLAYNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uploaded_model = aiplatform.Model.upload(\n",
    "    display_name=MODEL_DISPLAYNAME,\n",
    "    artifact_uri=f\"gs://{BUCKET}/{MODEL_DISPLAYNAME}\",\n",
    "    serving_container_image_uri=SERVING_CONTAINER_IMAGE_URI,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MACHINE_TYPE = \"n1-standard-4\"\n",
    "\n",
    "endpoint = uploaded_model.deploy(\n",
    "    machine_type=MACHINE_TYPE,\n",
    "    accelerator_type=None,\n",
    "    accelerator_count=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the model has been uploaded to the endpoint, you can query the endpoint to get predictions from the model. The query has to be a list of instances. From the signature of the model we see that the key for the query has to be `text`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instances = [\n",
    "    {\"text\": [\"I loved the movie and highly recomment it\"]},\n",
    "    {\"text\": [\"It was an okay movie in my opinion\"]},\n",
    "    {\"text\": [\"I hated the movie\"]},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = endpoint.predict(instances=instances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\" prediction:\", response.predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Cleanup\n",
    "When deploying a model to an endpoint for online prediction, the minimum min-replica-count is 1, and it is charged per node hour. So let's delete the endpoint to reduce unnecessary charges. Before we can delete the endpoint, we first undeploy all attached models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint.undeploy_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...then delete the endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint.delete()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Continued Learning\n",
    "\n",
    "In this lab, we chose small BERT to train our text classifier. There are other pre-trained BERT models which you can find here. Consider experiementing with some of these. However, remember that the bigger the model you choose to fine-tune, the longer it will take to train.\n",
    "\n",
    "There are \n",
    "  - [BERT-Base](https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3), [Uncased](https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3) and [seven more models](https://tfhub.dev/google/collections/bert/1) with trained weights released by the original BERT authors.\n",
    "  - [Small BERTs](https://tfhub.dev/google/collections/bert/1) have the same general architecture but fewer and/or smaller Transformer blocks, which lets you explore tradeoffs between speed, size and quality.\n",
    "  - [ALBERT](https://tfhub.dev/google/collections/albert/1): four different sizes of \"A Lite BERT\" that reduces model size (but not computation time) by sharing parameters between layers.\n",
    "  - [BERT Experts](https://tfhub.dev/google/collections/experts/bert/1): eight models that all have the BERT-base architecture but offer a choice between different pre-training domains, to align more closely with the target task.\n",
    "  - [Electra](https://tfhub.dev/google/collections/electra/1) has the same architecture as BERT (in three different sizes), but gets pre-trained as a discriminator in a set-up that resembles a Generative Adversarial Network (GAN).\n",
    "  - BERT with Talking-Heads Attention and Gated GELU [[base](https://tfhub.dev/tensorflow/talkheads_ggelu_bert_en_base/1), [large](https://tfhub.dev/tensorflow/talkheads_ggelu_bert_en_large/1)] has two improvements to the core of the Transformer architecture.\n",
    "\n",
    "The model documentation on TensorFlow Hub has more details and references to the\n",
    "research literature.\n",
    "\n",
    "Aside from the models available above, there are [multiple versions](https://tfhub.dev/google/collections/transformer_encoders_text/1) of the models that are larger and can yeld even better accuracy but they are too big to be fine-tuned on a single GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### License\n",
    "\n",
    "Copyright 2022 Google Inc. 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 http://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"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "classify_text_with_bert.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "bert_kernel",
   "name": "tf2-gpu.2-8.m106",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-8:m106"
  },
  "kernelspec": {
   "display_name": "bert_kernel",
   "language": "python",
   "name": "bert_kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}