From 9e912e76799e682370b2af28a234058eab7fe4e0 Mon Sep 17 00:00:00 2001
From: Susnato Dhar <susnatodhar10@gmail.com>
Date: Thu, 23 Feb 2023 12:57:07 +0530
Subject: [PATCH 1/5] bert_tiny_uncased_en_sst2 added

---
 ...rt_tiny_uncased_en_sst2_training_new.ipynb | 401 ++++++++++++++++++
 1 file changed, 401 insertions(+)
 create mode 100644 tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb

diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb
new file mode 100644
index 0000000000..71bfbb049d
--- /dev/null
+++ b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb
@@ -0,0 +1,401 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "gpuClass": "standard",
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# keras-nlp installation"
+      ],
+      "metadata": {
+        "id": "FKTVkreu3MGG"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "NtcTiJqG20LV",
+        "outputId": "d03cd2ad-50ea-4553-8c91-255bcfed9065"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
+          ]
+        }
+      ],
+      "source": [
+        "!pip install -q git+https://github.com/keras-team/keras-nlp.git tensorflow --upgrade"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Imports"
+      ],
+      "metadata": {
+        "id": "ofxRv3CK3Pdp"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "from tensorflow import keras\n",
+        "import tensorflow_datasets as tfds"
+      ],
+      "metadata": {
+        "id": "AHJ-nbcT3ATU"
+      },
+      "execution_count": 2,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Data"
+      ],
+      "metadata": {
+        "id": "r86tfpfB3T6_"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Load Dataset\n",
+        "\n",
+        "train_ds, valid_ds = tfds.load(\n",
+        "    \"glue/sst2\",\n",
+        "    split=[\"train\", \"validation\"],\n",
+        "    batch_size=16,\n",
+        ")\n",
+        "\n",
+        "def split_features(x):\n",
+        "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
+        "    # (features, label), where features is a tuple consisting of all\n",
+        "    # features.\n",
+        "    features = x[\"sentence\"]\n",
+        "    label = x[\"label\"]\n",
+        "    return (features, label)\n",
+        "\n",
+        "train_ds = train_ds.map(\n",
+        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+        ").prefetch(tf.data.AUTOTUNE)\n",
+        "valid_ds = valid_ds.map(\n",
+        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+        ").prefetch(tf.data.AUTOTUNE)"
+      ],
+      "metadata": {
+        "id": "8UDVK92k3VKQ"
+      },
+      "execution_count": 3,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Look first training set batch\n",
+        "# The format is (string_tensor, label_tensor)\n",
+        "train_ds.take(1).get_single_element()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "w8cfKvOl3aiT",
+        "outputId": "25c58ee3-88b2-4c07-9973-6bc0b966c875"
+      },
+      "execution_count": 4,
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
+              " array([b'for the uninitiated plays better on video with the sound ',\n",
+              "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
+              "        b'company once again dazzle and delight us ',\n",
+              "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
+              "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
+              "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
+              "        b'only masochistic moviegoers need apply . ',\n",
+              "        b'convince almost everyone that it was put on the screen , just for them ',\n",
+              "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
+              "        b'his supple understanding of the role ',\n",
+              "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
+              "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
+              "        b'of the startling intimacy ', b'draggy ',\n",
+              "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
+              "       dtype=object)>,\n",
+              " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 4
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Model"
+      ],
+      "metadata": {
+        "id": "CZTYpu0y4Y6A"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Create the Classifier Model \n",
+        "# For more details please look https://keras.io/guides/keras_nlp/getting_started/\n",
+        "\n",
+        "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
+        "    \"bert_tiny_en_uncased\",\n",
+        "    num_classes=2,\n",
+        "    dropout=0.1\n",
+        ")\n",
+        "\n",
+        "# Add loss function, optimizer and metrics\n",
+        "classifier.compile(\n",
+        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
+        "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
+        "    jit_compile=True,\n",
+        ")\n",
+        "   \n",
+        "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need \n",
+        "# to do it if we just want to train it on a downstream task.\n",
+        "classifier.layers[-1].build([None, ])\n",
+        "classifier.summary()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "1KxlC7gY32MH",
+        "outputId": "b02c8c90-9f37-40b8-843b-5fcc66819e80"
+      },
+      "execution_count": 5,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Model: \"bert_classifier\"\n",
+            "__________________________________________________________________________________________________\n",
+            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+            "==================================================================================================\n",
+            " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " bert_backbone (BertBackbone)   {'sequence_output':  4385920     ['padding_mask[0][0]',           \n",
+            "                                 (None, None, 128),               'segment_ids[0][0]',            \n",
+            "                                 'pooled_output': (               'token_ids[0][0]']              \n",
+            "                                None, 128)}                                                       \n",
+            "                                                                                                  \n",
+            " dropout (Dropout)              (None, 128)          0           ['bert_backbone[0][0]']          \n",
+            "                                                                                                  \n",
+            " logits (Dense)                 (None, 2)            258         ['dropout[0][0]']                \n",
+            "                                                                                                  \n",
+            " bert_preprocessor (BertPreproc  multiple            0           []                               \n",
+            " essor)                                                                                           \n",
+            "                                                                                                  \n",
+            "==================================================================================================\n",
+            "Total params: 4,386,178\n",
+            "Trainable params: 4,386,178\n",
+            "Non-trainable params: 0\n",
+            "__________________________________________________________________________________________________\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Train the model\n",
+        "\n",
+        "N_EPOCHS = 2\n",
+        "classifier.fit(\n",
+        "    train_ds,\n",
+        "    validation_data=valid_ds,\n",
+        "    epochs=N_EPOCHS,\n",
+        ")\n",
+        "  "
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "fxVfW7t14wRS",
+        "outputId": "b9a4277d-4fa1-4668-9373-d8cef77bcc6a"
+      },
+      "execution_count": 6,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
+            "Instructions for updating:\n",
+            "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Epoch 1/2\n",
+            "4210/4210 [==============================] - 828s 191ms/step - loss: 0.3782 - sparse_categorical_accuracy: 0.8299 - val_loss: 0.4344 - val_sparse_categorical_accuracy: 0.8165\n",
+            "Epoch 2/2\n",
+            "4210/4210 [==============================] - 783s 186ms/step - loss: 0.2409 - sparse_categorical_accuracy: 0.9039 - val_loss: 0.4626 - val_sparse_categorical_accuracy: 0.8222\n"
+          ]
+        },
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<keras.callbacks.History at 0x7f2390120df0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 6
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Save the weights\n",
+        "\n",
+        "model_name = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
+        "classifier.save_weights(f\"/content/{model_name}\")"
+      ],
+      "metadata": {
+        "id": "QExT3HYu6g6h"
+      },
+      "execution_count": 7,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Load the weigths\n",
+        "\n",
+        "weights_path = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
+        "classifier.load_weights(weights_path)"
+      ],
+      "metadata": {
+        "id": "8Y0ZTKRj62nq"
+      },
+      "execution_count": 8,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Test the model"
+      ],
+      "metadata": {
+        "id": "XWUIUhw47drP"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# We will shuffle the valid dataset and take 1st example\n",
+        "\n",
+        "shuffled_valid_ds = valid_ds.shuffle(55).rebatch(1)\n",
+        "element = shuffled_valid_ds.take(1).get_single_element()\n",
+        "\n",
+        "pred_logits = classifier.predict(element[0])\n",
+        "pred = tf.argmax(pred_logits)\n",
+        "\n",
+        "print(f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\")"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "FoQk4cy27jSk",
+        "outputId": "aff204c9-4879-403b-91ce-afcc3679b103"
+      },
+      "execution_count": 9,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "1/1 [==============================] - 3s 3s/step\n",
+            "Text :: preaches to two completely different choirs at the same time , which is a pretty amazing accomplishment .  \n",
+            "Label :: 1 \n",
+            "Model Prediction :: 1\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# You can test the model with your own statement!\n",
+        "\n",
+        "label_dict = {0:\"Bad Statement\", 1:\"Good Statement\"} # As for this dataset\n",
+        "\n",
+        "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
+        "output = tf.argmax(classifier.predict([\"Lord of the rings is best\"]), axis=1).numpy()[0]\n",
+        "print(label_dict[output])"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "7MB2U6Bh_j_m",
+        "outputId": "37a8223a-4178-4176-e660-258ed9b856c7"
+      },
+      "execution_count": 10,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "1/1 [==============================] - 1s 902ms/step\n",
+            "Good Statement\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "N-uuvlp2_ybw"
+      },
+      "execution_count": 10,
+      "outputs": []
+    }
+  ]
+}
\ No newline at end of file

From ec8a2f79635e83664c8ed9c31be3a83d22ad96e3 Mon Sep 17 00:00:00 2001
From: Susnato Dhar <susnatodhar10@gmail.com>
Date: Sun, 5 Mar 2023 17:02:00 +0530
Subject: [PATCH 2/5] replaced the file

---
 .../bert_tiny_uncased_en_sst2_training.ipynb  | 4364 ++---------------
 ...rt_tiny_uncased_en_sst2_training_new.ipynb |  401 --
 2 files changed, 352 insertions(+), 4413 deletions(-)
 delete mode 100644 tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb

diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
index 623a4ab83f..71bfbb049d 100644
--- a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
+++ b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
@@ -1,4061 +1,401 @@
 {
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "colab_type": "text",
-    "id": "view-in-github"
-   },
-   "source": [
-    "<a href=\"https://colab.research.google.com/github/jbischof/keras-nlp/blob/classifier_preset/bert_tiny_en_uncased_sst2_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
     "colab": {
-     "base_uri": "https://localhost:8080/"
+      "provenance": []
     },
-    "id": "2h__HbR944VZ",
-    "outputId": "ebbee015-5caf-453b-d246-42ce6f2d846d",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[K     |████████████████████████████████| 578.0 MB 16 kB/s \n",
-      "\u001b[K     |████████████████████████████████| 5.8 MB 72.0 MB/s \n",
-      "\u001b[K     |████████████████████████████████| 5.9 MB 68.7 MB/s \n",
-      "\u001b[K     |████████████████████████████████| 438 kB 88.6 MB/s \n",
-      "\u001b[K     |████████████████████████████████| 1.7 MB 80.4 MB/s \n",
-      "\u001b[K     |████████████████████████████████| 5.9 MB 61.9 MB/s \n",
-      "\u001b[?25h  Building wheel for keras-nlp (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
-     ]
-    }
-   ],
-   "source": [
-    "!pip install -q git+https://github.com/mattdangerw/keras-nlp.git@bert-pipeline tensorflow==2.10 --upgrade"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "rlmPYqUn7YJm",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "import keras_nlp\n",
-    "import tensorflow as tf\n",
-    "from tensorflow import keras\n",
-    "import tensorflow_datasets as tfds"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "_hE-veD7utTT"
-   },
-   "source": [
-    "# Data"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/",
-     "height": 498,
-     "referenced_widgets": [
-      "7dac43c361074f0f998d8c150aa922b9",
-      "02ecb50a721f450b940172d3d35a4abe",
-      "3977c5ac1af645b0b2c1e5886ac037ad",
-      "e9a881d7fcb84bafb6a063009bf4dd54",
-      "efdf1618f9224961b840bf58fb763876",
-      "071de690987f46d5b11d30fd416d0903",
-      "75083b32cd854cdd9de6490b7f828140",
-      "964b0fedfdbe4b4a990d3b72d68b4112",
-      "2308d8d3d7064c62ae470850e9b06a29",
-      "77c0fa586d4d4dcd982513c827799adb",
-      "a70a3e92a37247eeaad91ef1e140aa54",
-      "6bed029127e44fefa0ccd576edb855e9",
-      "289e60ba017345b1b61adeecc24bc2cb",
-      "bb9d3ced705148b6a8c3b71498bd9d81",
-      "4579a6614cd64cdcab0ef6c615cad8bf",
-      "5b7a72d055cd4cb68a4a9706bd276e1f",
-      "25e82eaf0b9849ada515c712b359b21d",
-      "1aca96855cbe4b41a050bdb59f97a903",
-      "078e2f25095d40cc9a96ad3488030022",
-      "16dc3ec494f6456794b43d33a95ef042",
-      "8c3528b24aba4002b71168fd0b8beb44",
-      "d9d4d67205d340a8acb919b7c80a842a",
-      "71a33b25dece4caca90a91160ea2382f",
-      "c84704566a0a45249c6973c8e57e3856",
-      "44ff2d0794b74097af53a8116da941a7",
-      "b22d3db2d68046a1afbca22b13ab7152",
-      "acbe1228b9bb42da9fc19e57e4dc6e74",
-      "06fa78192fd74020962d609ac59af6d5",
-      "4de362ac37224f968d2f6a4f2336b327",
-      "fac49204e83e49c9920c29a61382d9b7",
-      "8bf597ef7f7a4ce1a66cbfc1aa74a216",
-      "b96bbee1b3254bc397828295b8b6528a",
-      "4d7f2052da7f43f2a41ebfaaf562ae77",
-      "c96bafeb19454961aa95b819bec135df",
-      "879fa9af242b4c698c15747315a91ba9",
-      "cda35a02449442e48f3b6131547e1628",
-      "9b65c0c17d3d49589feed767a4ced04c",
-      "09d1f164a0e941438f06a3feb58bdd30",
-      "27d5831d71b547c081cb2dec9a42b300",
-      "14df05ca9bc642d9ad588db644ab7cf7",
-      "8fe22f392c204fdeb8c8b34b9c4bdc12",
-      "97717222d177460bad6c626004f3dd9d",
-      "78d934805a364c8989efba79bead117d",
-      "5dd04d5d68a54166b432bdf53dad0792",
-      "96d397cf27714dca8a8545c3ceb13d73",
-      "91a8aba5cca943fcb2f3f3353450aa78",
-      "10e9f95c33dd4b7b88150d6a7825e973",
-      "d929b9d20b1a4b00bb5902c0883b4034",
-      "c1976cfcd5c4438bb769ad5a77b0b265",
-      "42117d78043a4336bc9c90c88741c15a",
-      "b28ba82c231545e0b509a93ba0470449",
-      "d410bad163c24c1aa5e00d54090b05db",
-      "9eb1c4b7b7e94f8599705b2d1d509497",
-      "e9113eab130140f9b9367a61b1e4b5c0",
-      "454dfc6558ae4a1b83621ea1352fc051",
-      "1f5f58f9400b48a991a26786568c75af",
-      "d9a622959c0d4c45b0b1932d7c85893d",
-      "6e2dd84de8b3483e9f18142fa30177eb",
-      "6dc9fbdebed24e35bccaf0e92a6a491e",
-      "d147c549823e407eb843750f0ec45cbc",
-      "758811c0d8284c76984fc2ac974d835a",
-      "4e06b19ebe63441ebb6a9b1cd57e2d21",
-      "69a3fcf1ec174cbaa11f93fa1dded564",
-      "2522eb90704c4fd6be3a0223dd16442e",
-      "b0999cce50384ccfa1fb6dedcf6bc1d3",
-      "7f65a8e8ab6a411286d11c0125b316d0",
-      "cd67f31d575f44d9944a5019b3e38c19",
-      "cb1357c15a4c4a28a6198a61af82848b",
-      "0a9d89c682624b208b399ea9d8c26b0b",
-      "130007f46b794a4c9d283aa5fb333d04",
-      "a7fb5284dab64287b108a848c7203dd5",
-      "e70cb8396ab341c1bb0d60d2ebdd708a",
-      "d5020ccf7aff400988f0f0e6d1ff3557",
-      "8a0f69f4dc4040559001183e68b73aff",
-      "6f7051dbe4ad467bada9e50157141145",
-      "90e15c8359274c5a85a7cf427cd9ca22",
-      "b3086aaf2d8b491fbb5569d6226bfd29",
-      "52fca5d8bd9b4229bd64ebce81043160",
-      "8388b866783e4d9eb11093a99eaa3116",
-      "5df67222f0f74e1c9b57a59dbb1d3d98",
-      "04a3c2356d10466fb1683d202208f66d",
-      "335fe992357d4ca298312a55466c1d41",
-      "e5d9c734fb3e412aabc03e5790b12382",
-      "07fd35fa2e634720b064998b96e0e60e",
-      "ba64c4461ad14990afe4c1df45ad8d7d",
-      "30b1448d901a4191ab5bd62e89028e75",
-      "cbd21b3e1e5f4e318940805ac4d142ec",
-      "20ffd17f760c45788040797926b7fc35",
-      "a0d66c335f174159a7b2df88cb4b607f",
-      "58dfc96028a948c7af85f4df518bc171",
-      "436d985f48104913933079fdb44fde6f",
-      "aca24255c81943aeb92a036311e53cef",
-      "c7a31a4a0e4d4b27a93e48950acfe5e5",
-      "b65aff687d4242fa94ede609a48ad847",
-      "1fc728a0824144d1a60399e3d8529ea9",
-      "5634ff298df447dd92f4b34cf382b7e0",
-      "4b2ed16c315e413188360923ec7e19eb",
-      "80053cb69a2a4b2c8ad3c038fc3e2c39",
-      "7c9acf3891a040f5af892b86e72c803c",
-      "b4a60897f73642eeb692cafa7997167c",
-      "c4d8f2f2b9604745ac03fe5f4da7eeee",
-      "87e1ae4c68394e23ba0ed21f8fffb5a6",
-      "670f7df971d14fabbffcfd2535e90005",
-      "148f6232f6db44919f04a78c70aba8c1",
-      "66071ddc16014e9a8604f37fe1441b33",
-      "88b26fe32ba64f9e90c10b6c40381635",
-      "4679a29c8e2c4fa686713038e1c4686e",
-      "b9d49ffe543342c694905d2dee72b880",
-      "93c4b10a8f074c2c95cee65850ab2040",
-      "275c1f7417e247ce99519f0ef43738c9"
-     ]
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
     },
-    "id": "e3Q1wpT6mh5o",
-    "outputId": "2537d487-6dc3-4072-a03d-0e5f1879e1e1",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Downloading and preparing dataset 7.09 MiB (download: 7.09 MiB, generated: 7.22 MiB, total: 14.31 MiB) to ~/tensorflow_datasets/glue/sst2/2.0.0...\n"
-     ]
+    "language_info": {
+      "name": "python"
     },
+    "gpuClass": "standard",
+    "accelerator": "GPU"
+  },
+  "cells": [
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7dac43c361074f0f998d8c150aa922b9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Dl Completed...: 0 url [00:00, ? url/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+      "cell_type": "markdown",
+      "source": [
+        "# keras-nlp installation"
+      ],
+      "metadata": {
+        "id": "FKTVkreu3MGG"
+      }
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6bed029127e44fefa0ccd576edb855e9",
-       "version_major": 2,
-       "version_minor": 0
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "NtcTiJqG20LV",
+        "outputId": "d03cd2ad-50ea-4553-8c91-255bcfed9065"
       },
-      "text/plain": [
-       "Dl Size...: 0 MiB [00:00, ? MiB/s]"
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
+          ]
+        }
+      ],
+      "source": [
+        "!pip install -q git+https://github.com/keras-team/keras-nlp.git tensorflow --upgrade"
       ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "71a33b25dece4caca90a91160ea2382f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Extraction completed...: 0 file [00:00, ? file/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+      "cell_type": "markdown",
+      "source": [
+        "# Imports"
+      ],
+      "metadata": {
+        "id": "ofxRv3CK3Pdp"
+      }
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c96bafeb19454961aa95b819bec135df",
-       "version_major": 2,
-       "version_minor": 0
+      "cell_type": "code",
+      "source": [
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "from tensorflow import keras\n",
+        "import tensorflow_datasets as tfds"
+      ],
+      "metadata": {
+        "id": "AHJ-nbcT3ATU"
       },
-      "text/plain": [
-       "Generating splits...:   0%|          | 0/3 [00:00<?, ? splits/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+      "execution_count": 2,
+      "outputs": []
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "96d397cf27714dca8a8545c3ceb13d73",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Generating train examples...:   0%|          | 0/67349 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+      "cell_type": "markdown",
+      "source": [
+        "#Data"
+      ],
+      "metadata": {
+        "id": "r86tfpfB3T6_"
+      }
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1f5f58f9400b48a991a26786568c75af",
-       "version_major": 2,
-       "version_minor": 0
+      "cell_type": "code",
+      "source": [
+        "# Load Dataset\n",
+        "\n",
+        "train_ds, valid_ds = tfds.load(\n",
+        "    \"glue/sst2\",\n",
+        "    split=[\"train\", \"validation\"],\n",
+        "    batch_size=16,\n",
+        ")\n",
+        "\n",
+        "def split_features(x):\n",
+        "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
+        "    # (features, label), where features is a tuple consisting of all\n",
+        "    # features.\n",
+        "    features = x[\"sentence\"]\n",
+        "    label = x[\"label\"]\n",
+        "    return (features, label)\n",
+        "\n",
+        "train_ds = train_ds.map(\n",
+        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+        ").prefetch(tf.data.AUTOTUNE)\n",
+        "valid_ds = valid_ds.map(\n",
+        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+        ").prefetch(tf.data.AUTOTUNE)"
+      ],
+      "metadata": {
+        "id": "8UDVK92k3VKQ"
       },
-      "text/plain": [
-       "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-train.tfrecord*...:   0%|          | 0/6…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+      "execution_count": 3,
+      "outputs": []
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd67f31d575f44d9944a5019b3e38c19",
-       "version_major": 2,
-       "version_minor": 0
+      "cell_type": "code",
+      "source": [
+        "# Look first training set batch\n",
+        "# The format is (string_tensor, label_tensor)\n",
+        "train_ds.take(1).get_single_element()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "w8cfKvOl3aiT",
+        "outputId": "25c58ee3-88b2-4c07-9973-6bc0b966c875"
       },
-      "text/plain": [
-       "Generating validation examples...:   0%|          | 0/872 [00:00<?, ? examples/s]"
+      "execution_count": 4,
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
+              " array([b'for the uninitiated plays better on video with the sound ',\n",
+              "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
+              "        b'company once again dazzle and delight us ',\n",
+              "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
+              "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
+              "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
+              "        b'only masochistic moviegoers need apply . ',\n",
+              "        b'convince almost everyone that it was put on the screen , just for them ',\n",
+              "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
+              "        b'his supple understanding of the role ',\n",
+              "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
+              "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
+              "        b'of the startling intimacy ', b'draggy ',\n",
+              "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
+              "       dtype=object)>,\n",
+              " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 4
+        }
       ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "52fca5d8bd9b4229bd64ebce81043160",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-validation.tfrecord*...:   0%|          …"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+      "cell_type": "markdown",
+      "source": [
+        "#Model"
+      ],
+      "metadata": {
+        "id": "CZTYpu0y4Y6A"
+      }
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a0d66c335f174159a7b2df88cb4b607f",
-       "version_major": 2,
-       "version_minor": 0
+      "cell_type": "code",
+      "source": [
+        "# Create the Classifier Model \n",
+        "# For more details please look https://keras.io/guides/keras_nlp/getting_started/\n",
+        "\n",
+        "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
+        "    \"bert_tiny_en_uncased\",\n",
+        "    num_classes=2,\n",
+        "    dropout=0.1\n",
+        ")\n",
+        "\n",
+        "# Add loss function, optimizer and metrics\n",
+        "classifier.compile(\n",
+        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
+        "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
+        "    jit_compile=True,\n",
+        ")\n",
+        "   \n",
+        "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need \n",
+        "# to do it if we just want to train it on a downstream task.\n",
+        "classifier.layers[-1].build([None, ])\n",
+        "classifier.summary()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "1KxlC7gY32MH",
+        "outputId": "b02c8c90-9f37-40b8-843b-5fcc66819e80"
       },
-      "text/plain": [
-       "Generating test examples...:   0%|          | 0/1821 [00:00<?, ? examples/s]"
+      "execution_count": 5,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Model: \"bert_classifier\"\n",
+            "__________________________________________________________________________________________________\n",
+            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+            "==================================================================================================\n",
+            " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " bert_backbone (BertBackbone)   {'sequence_output':  4385920     ['padding_mask[0][0]',           \n",
+            "                                 (None, None, 128),               'segment_ids[0][0]',            \n",
+            "                                 'pooled_output': (               'token_ids[0][0]']              \n",
+            "                                None, 128)}                                                       \n",
+            "                                                                                                  \n",
+            " dropout (Dropout)              (None, 128)          0           ['bert_backbone[0][0]']          \n",
+            "                                                                                                  \n",
+            " logits (Dense)                 (None, 2)            258         ['dropout[0][0]']                \n",
+            "                                                                                                  \n",
+            " bert_preprocessor (BertPreproc  multiple            0           []                               \n",
+            " essor)                                                                                           \n",
+            "                                                                                                  \n",
+            "==================================================================================================\n",
+            "Total params: 4,386,178\n",
+            "Trainable params: 4,386,178\n",
+            "Non-trainable params: 0\n",
+            "__________________________________________________________________________________________________\n"
+          ]
+        }
       ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b4a60897f73642eeb692cafa7997167c",
-       "version_major": 2,
-       "version_minor": 0
+      "cell_type": "code",
+      "source": [
+        "# Train the model\n",
+        "\n",
+        "N_EPOCHS = 2\n",
+        "classifier.fit(\n",
+        "    train_ds,\n",
+        "    validation_data=valid_ds,\n",
+        "    epochs=N_EPOCHS,\n",
+        ")\n",
+        "  "
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "fxVfW7t14wRS",
+        "outputId": "b9a4277d-4fa1-4668-9373-d8cef77bcc6a"
       },
-      "text/plain": [
-       "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-test.tfrecord*...:   0%|          | 0/18…"
+      "execution_count": 6,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
+            "Instructions for updating:\n",
+            "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Epoch 1/2\n",
+            "4210/4210 [==============================] - 828s 191ms/step - loss: 0.3782 - sparse_categorical_accuracy: 0.8299 - val_loss: 0.4344 - val_sparse_categorical_accuracy: 0.8165\n",
+            "Epoch 2/2\n",
+            "4210/4210 [==============================] - 783s 186ms/step - loss: 0.2409 - sparse_categorical_accuracy: 0.9039 - val_loss: 0.4626 - val_sparse_categorical_accuracy: 0.8222\n"
+          ]
+        },
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<keras.callbacks.History at 0x7f2390120df0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 6
+        }
       ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
     },
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Dataset glue downloaded and prepared to ~/tensorflow_datasets/glue/sst2/2.0.0. Subsequent calls will reuse this data.\n"
-     ]
+      "cell_type": "code",
+      "source": [
+        "# Save the weights\n",
+        "\n",
+        "model_name = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
+        "classifier.save_weights(f\"/content/{model_name}\")"
+      ],
+      "metadata": {
+        "id": "QExT3HYu6g6h"
+      },
+      "execution_count": 7,
+      "outputs": []
     },
     {
-     "data": {
-      "text/plain": [
-       "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
-       " array([b'for the uninitiated plays better on video with the sound ',\n",
-       "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
-       "        b'company once again dazzle and delight us ',\n",
-       "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
-       "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
-       "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
-       "        b'only masochistic moviegoers need apply . ',\n",
-       "        b'convince almost everyone that it was put on the screen , just for them ',\n",
-       "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
-       "        b'his supple understanding of the role ',\n",
-       "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
-       "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
-       "        b'of the startling intimacy ', b'draggy ',\n",
-       "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
-       "       dtype=object)>,\n",
-       " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "train_ds, valid_ds = tfds.load(\n",
-    "    \"glue/sst2\",\n",
-    "    split=[\"train\", \"validation\"],\n",
-    "    batch_size=16,\n",
-    ")\n",
-    "\n",
-    "\n",
-    "def split_features(x):\n",
-    "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
-    "    # (features, label), where features is a tuple consisting of all\n",
-    "    # features.\n",
-    "    features = x[\"sentence\"]\n",
-    "    label = x[\"label\"]\n",
-    "    return (features, label)\n",
-    "\n",
-    "\n",
-    "train_ds = train_ds.map(\n",
-    "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
-    ").prefetch(tf.data.AUTOTUNE)\n",
-    "valid_ds = valid_ds.map(\n",
-    "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
-    ").prefetch(tf.data.AUTOTUNE)\n",
-    "\n",
-    "# Look first training set batch\n",
-    "# The format is (string_tensor, label_tensor)\n",
-    "train_ds.take(1).get_single_element()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
+      "cell_type": "code",
+      "source": [
+        "# Load the weigths\n",
+        "\n",
+        "weights_path = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
+        "classifier.load_weights(weights_path)"
+      ],
+      "metadata": {
+        "id": "8Y0ZTKRj62nq"
+      },
+      "execution_count": 8,
+      "outputs": []
     },
-    "id": "_cj9caiYt_jb",
-    "outputId": "0f09fccb-97c5-44c0-c156-94f182bc5c05",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Epoch 1/2\n",
-      "4210/4210 [==============================] - 740s 173ms/step - loss: 0.3747 - sparse_categorical_accuracy: 0.8295 - val_loss: 0.4244 - val_sparse_categorical_accuracy: 0.8165\n",
-      "Epoch 2/2\n",
-      "4210/4210 [==============================] - 710s 169ms/step - loss: 0.2378 - sparse_categorical_accuracy: 0.9054 - val_loss: 0.4463 - val_sparse_categorical_accuracy: 0.8234\n"
-     ]
+      "cell_type": "markdown",
+      "source": [
+        "#Test the model"
+      ],
+      "metadata": {
+        "id": "XWUIUhw47drP"
+      }
     },
     {
-     "data": {
-      "text/plain": [
-       "<keras.callbacks.History at 0x7f6bbd0923d0>"
+      "cell_type": "code",
+      "source": [
+        "# We will shuffle the valid dataset and take 1st example\n",
+        "\n",
+        "shuffled_valid_ds = valid_ds.shuffle(55).rebatch(1)\n",
+        "element = shuffled_valid_ds.take(1).get_single_element()\n",
+        "\n",
+        "pred_logits = classifier.predict(element[0])\n",
+        "pred = tf.argmax(pred_logits)\n",
+        "\n",
+        "print(f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\")"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "FoQk4cy27jSk",
+        "outputId": "aff204c9-4879-403b-91ce-afcc3679b103"
+      },
+      "execution_count": 9,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "1/1 [==============================] - 3s 3s/step\n",
+            "Text :: preaches to two completely different choirs at the same time , which is a pretty amazing accomplishment .  \n",
+            "Label :: 1 \n",
+            "Model Prediction :: 1\n"
+          ]
+        }
       ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
-    "    \"bert_tiny_en_uncased\",\n",
-    "    num_classes=2,\n",
-    ")\n",
-    "classifier.compile(\n",
-    "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
-    "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
-    "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
-    "    jit_compile=True,\n",
-    ")\n",
-    "classifier.fit(\n",
-    "    train_ds,\n",
-    "    validation_data=valid_ds,\n",
-    "    epochs=2,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "fJcxn1si0dKV",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "classifier.save_weights(\"/content/model.h5\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "jFSJUTZEoUws",
-    "outputId": "93cd3576-0c42-43e0-8fd0-c4d9a42d3a87",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "-rw-r--r-- 1 root root 17M Nov 28 19:58 /content/model.h5\n",
-      "1f9c2d59f9e229e08f3fbd44239cfb0b  /content/model.h5\n"
-     ]
-    }
-   ],
-   "source": [
-    "!ls -lh /content/model.h5\n",
-    "!md5sum /content/model.h5"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "Wk6aVbYyUHBX",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "classifier.load_weights(\"/content/model.h5\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
     },
-    "id": "MJqmTfYtwt8V",
-    "outputId": "726b2a10-c650-4059-e15e-3173001ad444",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "{'backbone': {'class_name': 'keras_nlp>Bert',\n",
-       "  'config': {'vocabulary_size': 30522,\n",
-       "   'hidden_dim': 128,\n",
-       "   'intermediate_dim': 512,\n",
-       "   'num_layers': 2,\n",
-       "   'num_heads': 2,\n",
-       "   'max_sequence_length': 512,\n",
-       "   'num_segments': 2,\n",
-       "   'dropout': 0.1,\n",
-       "   'name': 'backbone',\n",
-       "   'trainable': True}},\n",
-       " 'preprocessor': {'class_name': 'keras_nlp>BertPreprocessor',\n",
-       "  'config': {'name': 'bert_preprocessor_1',\n",
-       "   'trainable': True,\n",
-       "   'dtype': 'float32',\n",
-       "   'tokenizer': {'class_name': 'keras_nlp>BertTokenizer',\n",
-       "    'config': {'name': 'bert_tokenizer_1',\n",
-       "     'trainable': True,\n",
-       "     'dtype': 'int32',\n",
-       "     'vocabulary': [],\n",
-       "     'sequence_length': None,\n",
-       "     'lowercase': True,\n",
-       "     'strip_accents': False,\n",
-       "     'split': True,\n",
-       "     'suffix_indicator': '##',\n",
-       "     'oov_token': '[UNK]'}},\n",
-       "   'sequence_length': 512,\n",
-       "   'truncate': 'round_robin'}},\n",
-       " 'num_classes': 2,\n",
-       " 'name': 'bert_classifier_1',\n",
-       " 'trainable': True}"
+      "cell_type": "code",
+      "source": [
+        "# You can test the model with your own statement!\n",
+        "\n",
+        "label_dict = {0:\"Bad Statement\", 1:\"Good Statement\"} # As for this dataset\n",
+        "\n",
+        "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
+        "output = tf.argmax(classifier.predict([\"Lord of the rings is best\"]), axis=1).numpy()[0]\n",
+        "print(label_dict[output])"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "7MB2U6Bh_j_m",
+        "outputId": "37a8223a-4178-4176-e660-258ed9b856c7"
+      },
+      "execution_count": 10,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "1/1 [==============================] - 1s 902ms/step\n",
+            "Good Statement\n"
+          ]
+        }
       ]
-     },
-     "execution_count": 23,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "config = classifier.get_config()\n",
-    "# Zero out vocab so readable\n",
-    "config[\"preprocessor\"][\"config\"][\"tokenizer\"][\"config\"][\"vocabulary\"] = []\n",
-    "config"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
     },
-    "id": "z113U3yutzY8",
-    "outputId": "763edfc4-fcd1-4216-9afb-276771258f6a",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
-   "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "{'backbone': {'class_name': 'keras_nlp>Bert',\n",
-       "  'config': {'vocabulary_size': 30522,\n",
-       "   'hidden_dim': 128,\n",
-       "   'intermediate_dim': 512,\n",
-       "   'num_layers': 2,\n",
-       "   'num_heads': 2,\n",
-       "   'max_sequence_length': 512,\n",
-       "   'num_segments': 2,\n",
-       "   'dropout': 0.1,\n",
-       "   'name': 'backbone',\n",
-       "   'trainable': True}},\n",
-       " 'num_classes': 2,\n",
-       " 'name': 'bert_classifier_1',\n",
-       " 'trainable': True}"
-      ]
-     },
-     "execution_count": 15,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "config.pop(\"preprocessor\")\n",
-    "config"
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "GPU",
-  "colab": {
-   "include_colab_link": true,
-   "provenance": []
-  },
-  "gpuClass": "premium",
-  "kernelspec": {
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "widgets": {
-   "application/vnd.jupyter.widget-state+json": {
-    "02ecb50a721f450b940172d3d35a4abe": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_071de690987f46d5b11d30fd416d0903",
-      "placeholder": "​",
-      "style": "IPY_MODEL_75083b32cd854cdd9de6490b7f828140",
-      "value": "Dl Completed...: 100%"
-     }
-    },
-    "04a3c2356d10466fb1683d202208f66d": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_cbd21b3e1e5f4e318940805ac4d142ec",
-      "placeholder": "​",
-      "style": "IPY_MODEL_20ffd17f760c45788040797926b7fc35",
-      "value": " 0/872 [00:00&lt;?, ? examples/s]"
-     }
-    },
-    "06fa78192fd74020962d609ac59af6d5": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "071de690987f46d5b11d30fd416d0903": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "078e2f25095d40cc9a96ad3488030022": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": "20px"
-     }
-    },
-    "07fd35fa2e634720b064998b96e0e60e": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "09d1f164a0e941438f06a3feb58bdd30": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "0a9d89c682624b208b399ea9d8c26b0b": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_8a0f69f4dc4040559001183e68b73aff",
-      "max": 872,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_6f7051dbe4ad467bada9e50157141145",
-      "value": 872
-     }
-    },
-    "10e9f95c33dd4b7b88150d6a7825e973": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_d410bad163c24c1aa5e00d54090b05db",
-      "max": 67349,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_9eb1c4b7b7e94f8599705b2d1d509497",
-      "value": 67349
-     }
-    },
-    "130007f46b794a4c9d283aa5fb333d04": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_90e15c8359274c5a85a7cf427cd9ca22",
-      "placeholder": "​",
-      "style": "IPY_MODEL_b3086aaf2d8b491fbb5569d6226bfd29",
-      "value": " 634/872 [00:00&lt;00:00, 6335.34 examples/s]"
-     }
-    },
-    "148f6232f6db44919f04a78c70aba8c1": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "14df05ca9bc642d9ad588db644ab7cf7": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "16dc3ec494f6456794b43d33a95ef042": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "1aca96855cbe4b41a050bdb59f97a903": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "1f5f58f9400b48a991a26786568c75af": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_d9a622959c0d4c45b0b1932d7c85893d",
-       "IPY_MODEL_6e2dd84de8b3483e9f18142fa30177eb",
-       "IPY_MODEL_6dc9fbdebed24e35bccaf0e92a6a491e"
-      ],
-      "layout": "IPY_MODEL_d147c549823e407eb843750f0ec45cbc"
-     }
-    },
-    "1fc728a0824144d1a60399e3d8529ea9": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "20ffd17f760c45788040797926b7fc35": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "2308d8d3d7064c62ae470850e9b06a29": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "2522eb90704c4fd6be3a0223dd16442e": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "25e82eaf0b9849ada515c712b359b21d": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "275c1f7417e247ce99519f0ef43738c9": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "27d5831d71b547c081cb2dec9a42b300": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "289e60ba017345b1b61adeecc24bc2cb": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_25e82eaf0b9849ada515c712b359b21d",
-      "placeholder": "​",
-      "style": "IPY_MODEL_1aca96855cbe4b41a050bdb59f97a903",
-      "value": "Dl Size...: 100%"
-     }
-    },
-    "30b1448d901a4191ab5bd62e89028e75": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "335fe992357d4ca298312a55466c1d41": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "3977c5ac1af645b0b2c1e5886ac037ad": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "success",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_964b0fedfdbe4b4a990d3b72d68b4112",
-      "max": 1,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_2308d8d3d7064c62ae470850e9b06a29",
-      "value": 1
-     }
-    },
-    "42117d78043a4336bc9c90c88741c15a": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "436d985f48104913933079fdb44fde6f": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_5634ff298df447dd92f4b34cf382b7e0",
-      "max": 1821,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_4b2ed16c315e413188360923ec7e19eb",
-      "value": 1821
-     }
-    },
-    "44ff2d0794b74097af53a8116da941a7": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "success",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_fac49204e83e49c9920c29a61382d9b7",
-      "max": 1,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_8bf597ef7f7a4ce1a66cbfc1aa74a216",
-      "value": 1
-     }
-    },
-    "454dfc6558ae4a1b83621ea1352fc051": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "4579a6614cd64cdcab0ef6c615cad8bf": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_8c3528b24aba4002b71168fd0b8beb44",
-      "placeholder": "​",
-      "style": "IPY_MODEL_d9d4d67205d340a8acb919b7c80a842a",
-      "value": " 7/7 [00:01&lt;00:00,  3.67 MiB/s]"
-     }
-    },
-    "4679a29c8e2c4fa686713038e1c4686e": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "4b2ed16c315e413188360923ec7e19eb": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "4d7f2052da7f43f2a41ebfaaf562ae77": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "4de362ac37224f968d2f6a4f2336b327": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "4e06b19ebe63441ebb6a9b1cd57e2d21": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "52fca5d8bd9b4229bd64ebce81043160": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_8388b866783e4d9eb11093a99eaa3116",
-       "IPY_MODEL_5df67222f0f74e1c9b57a59dbb1d3d98",
-       "IPY_MODEL_04a3c2356d10466fb1683d202208f66d"
-      ],
-      "layout": "IPY_MODEL_335fe992357d4ca298312a55466c1d41"
-     }
-    },
-    "5634ff298df447dd92f4b34cf382b7e0": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "58dfc96028a948c7af85f4df518bc171": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_b65aff687d4242fa94ede609a48ad847",
-      "placeholder": "​",
-      "style": "IPY_MODEL_1fc728a0824144d1a60399e3d8529ea9",
-      "value": "Generating test examples...:  77%"
-     }
-    },
-    "5b7a72d055cd4cb68a4a9706bd276e1f": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "5dd04d5d68a54166b432bdf53dad0792": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "5df67222f0f74e1c9b57a59dbb1d3d98": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_ba64c4461ad14990afe4c1df45ad8d7d",
-      "max": 872,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_30b1448d901a4191ab5bd62e89028e75",
-      "value": 872
-     }
-    },
-    "66071ddc16014e9a8604f37fe1441b33": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "670f7df971d14fabbffcfd2535e90005": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_93c4b10a8f074c2c95cee65850ab2040",
-      "placeholder": "​",
-      "style": "IPY_MODEL_275c1f7417e247ce99519f0ef43738c9",
-      "value": " 0/1821 [00:00&lt;?, ? examples/s]"
-     }
-    },
-    "69a3fcf1ec174cbaa11f93fa1dded564": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "6bed029127e44fefa0ccd576edb855e9": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_289e60ba017345b1b61adeecc24bc2cb",
-       "IPY_MODEL_bb9d3ced705148b6a8c3b71498bd9d81",
-       "IPY_MODEL_4579a6614cd64cdcab0ef6c615cad8bf"
-      ],
-      "layout": "IPY_MODEL_5b7a72d055cd4cb68a4a9706bd276e1f"
-     }
-    },
-    "6dc9fbdebed24e35bccaf0e92a6a491e": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_b0999cce50384ccfa1fb6dedcf6bc1d3",
-      "placeholder": "​",
-      "style": "IPY_MODEL_7f65a8e8ab6a411286d11c0125b316d0",
-      "value": " 52039/67349 [00:00&lt;00:00, 292448.11 examples/s]"
-     }
-    },
-    "6e2dd84de8b3483e9f18142fa30177eb": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_69a3fcf1ec174cbaa11f93fa1dded564",
-      "max": 67349,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_2522eb90704c4fd6be3a0223dd16442e",
-      "value": 67349
-     }
-    },
-    "6f7051dbe4ad467bada9e50157141145": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "71a33b25dece4caca90a91160ea2382f": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_c84704566a0a45249c6973c8e57e3856",
-       "IPY_MODEL_44ff2d0794b74097af53a8116da941a7",
-       "IPY_MODEL_b22d3db2d68046a1afbca22b13ab7152"
-      ],
-      "layout": "IPY_MODEL_acbe1228b9bb42da9fc19e57e4dc6e74"
-     }
-    },
-    "75083b32cd854cdd9de6490b7f828140": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "758811c0d8284c76984fc2ac974d835a": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "77c0fa586d4d4dcd982513c827799adb": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "78d934805a364c8989efba79bead117d": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "7c9acf3891a040f5af892b86e72c803c": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "7dac43c361074f0f998d8c150aa922b9": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_02ecb50a721f450b940172d3d35a4abe",
-       "IPY_MODEL_3977c5ac1af645b0b2c1e5886ac037ad",
-       "IPY_MODEL_e9a881d7fcb84bafb6a063009bf4dd54"
-      ],
-      "layout": "IPY_MODEL_efdf1618f9224961b840bf58fb763876"
-     }
-    },
-    "7f65a8e8ab6a411286d11c0125b316d0": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "80053cb69a2a4b2c8ad3c038fc3e2c39": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "8388b866783e4d9eb11093a99eaa3116": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_e5d9c734fb3e412aabc03e5790b12382",
-      "placeholder": "​",
-      "style": "IPY_MODEL_07fd35fa2e634720b064998b96e0e60e",
-      "value": "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-validation.tfrecord*...:   0%"
-     }
-    },
-    "879fa9af242b4c698c15747315a91ba9": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_27d5831d71b547c081cb2dec9a42b300",
-      "placeholder": "​",
-      "style": "IPY_MODEL_14df05ca9bc642d9ad588db644ab7cf7",
-      "value": "Generating splits...: 100%"
-     }
-    },
-    "87e1ae4c68394e23ba0ed21f8fffb5a6": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_4679a29c8e2c4fa686713038e1c4686e",
-      "max": 1821,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_b9d49ffe543342c694905d2dee72b880",
-      "value": 1821
-     }
-    },
-    "88b26fe32ba64f9e90c10b6c40381635": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "8a0f69f4dc4040559001183e68b73aff": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "8bf597ef7f7a4ce1a66cbfc1aa74a216": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "8c3528b24aba4002b71168fd0b8beb44": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "8fe22f392c204fdeb8c8b34b9c4bdc12": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "90e15c8359274c5a85a7cf427cd9ca22": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "91a8aba5cca943fcb2f3f3353450aa78": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_42117d78043a4336bc9c90c88741c15a",
-      "placeholder": "​",
-      "style": "IPY_MODEL_b28ba82c231545e0b509a93ba0470449",
-      "value": "Generating train examples...: 100%"
-     }
-    },
-    "93c4b10a8f074c2c95cee65850ab2040": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "964b0fedfdbe4b4a990d3b72d68b4112": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": "20px"
-     }
-    },
-    "96d397cf27714dca8a8545c3ceb13d73": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_91a8aba5cca943fcb2f3f3353450aa78",
-       "IPY_MODEL_10e9f95c33dd4b7b88150d6a7825e973",
-       "IPY_MODEL_d929b9d20b1a4b00bb5902c0883b4034"
-      ],
-      "layout": "IPY_MODEL_c1976cfcd5c4438bb769ad5a77b0b265"
-     }
-    },
-    "97717222d177460bad6c626004f3dd9d": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "9b65c0c17d3d49589feed767a4ced04c": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_78d934805a364c8989efba79bead117d",
-      "placeholder": "​",
-      "style": "IPY_MODEL_5dd04d5d68a54166b432bdf53dad0792",
-      "value": " 3/3 [00:09&lt;00:00,  2.22s/ splits]"
-     }
-    },
-    "9eb1c4b7b7e94f8599705b2d1d509497": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "a0d66c335f174159a7b2df88cb4b607f": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_58dfc96028a948c7af85f4df518bc171",
-       "IPY_MODEL_436d985f48104913933079fdb44fde6f",
-       "IPY_MODEL_aca24255c81943aeb92a036311e53cef"
-      ],
-      "layout": "IPY_MODEL_c7a31a4a0e4d4b27a93e48950acfe5e5"
-     }
-    },
-    "a70a3e92a37247eeaad91ef1e140aa54": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "a7fb5284dab64287b108a848c7203dd5": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "aca24255c81943aeb92a036311e53cef": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_80053cb69a2a4b2c8ad3c038fc3e2c39",
-      "placeholder": "​",
-      "style": "IPY_MODEL_7c9acf3891a040f5af892b86e72c803c",
-      "value": " 1402/1821 [00:00&lt;00:00, 7169.58 examples/s]"
-     }
-    },
-    "acbe1228b9bb42da9fc19e57e4dc6e74": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "b0999cce50384ccfa1fb6dedcf6bc1d3": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "b22d3db2d68046a1afbca22b13ab7152": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_b96bbee1b3254bc397828295b8b6528a",
-      "placeholder": "​",
-      "style": "IPY_MODEL_4d7f2052da7f43f2a41ebfaaf562ae77",
-      "value": " 1/1 [00:01&lt;00:00,  1.28s/ file]"
-     }
-    },
-    "b28ba82c231545e0b509a93ba0470449": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "b3086aaf2d8b491fbb5569d6226bfd29": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "b4a60897f73642eeb692cafa7997167c": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_c4d8f2f2b9604745ac03fe5f4da7eeee",
-       "IPY_MODEL_87e1ae4c68394e23ba0ed21f8fffb5a6",
-       "IPY_MODEL_670f7df971d14fabbffcfd2535e90005"
-      ],
-      "layout": "IPY_MODEL_148f6232f6db44919f04a78c70aba8c1"
-     }
-    },
-    "b65aff687d4242fa94ede609a48ad847": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "b96bbee1b3254bc397828295b8b6528a": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "b9d49ffe543342c694905d2dee72b880": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": ""
-     }
-    },
-    "ba64c4461ad14990afe4c1df45ad8d7d": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "bb9d3ced705148b6a8c3b71498bd9d81": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "success",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_078e2f25095d40cc9a96ad3488030022",
-      "max": 1,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_16dc3ec494f6456794b43d33a95ef042",
-      "value": 1
-     }
-    },
-    "c1976cfcd5c4438bb769ad5a77b0b265": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "c4d8f2f2b9604745ac03fe5f4da7eeee": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_66071ddc16014e9a8604f37fe1441b33",
-      "placeholder": "​",
-      "style": "IPY_MODEL_88b26fe32ba64f9e90c10b6c40381635",
-      "value": "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-test.tfrecord*...:   0%"
-     }
-    },
-    "c7a31a4a0e4d4b27a93e48950acfe5e5": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "c84704566a0a45249c6973c8e57e3856": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_06fa78192fd74020962d609ac59af6d5",
-      "placeholder": "​",
-      "style": "IPY_MODEL_4de362ac37224f968d2f6a4f2336b327",
-      "value": "Extraction completed...: 100%"
-     }
-    },
-    "c96bafeb19454961aa95b819bec135df": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_879fa9af242b4c698c15747315a91ba9",
-       "IPY_MODEL_cda35a02449442e48f3b6131547e1628",
-       "IPY_MODEL_9b65c0c17d3d49589feed767a4ced04c"
-      ],
-      "layout": "IPY_MODEL_09d1f164a0e941438f06a3feb58bdd30"
-     }
-    },
-    "cb1357c15a4c4a28a6198a61af82848b": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_e70cb8396ab341c1bb0d60d2ebdd708a",
-      "placeholder": "​",
-      "style": "IPY_MODEL_d5020ccf7aff400988f0f0e6d1ff3557",
-      "value": "Generating validation examples...:  73%"
-     }
-    },
-    "cbd21b3e1e5f4e318940805ac4d142ec": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "cd67f31d575f44d9944a5019b3e38c19": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_cb1357c15a4c4a28a6198a61af82848b",
-       "IPY_MODEL_0a9d89c682624b208b399ea9d8c26b0b",
-       "IPY_MODEL_130007f46b794a4c9d283aa5fb333d04"
-      ],
-      "layout": "IPY_MODEL_a7fb5284dab64287b108a848c7203dd5"
-     }
-    },
-    "cda35a02449442e48f3b6131547e1628": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_8fe22f392c204fdeb8c8b34b9c4bdc12",
-      "max": 3,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_97717222d177460bad6c626004f3dd9d",
-      "value": 3
-     }
-    },
-    "d147c549823e407eb843750f0ec45cbc": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": "hidden",
-      "width": null
-     }
-    },
-    "d410bad163c24c1aa5e00d54090b05db": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "d5020ccf7aff400988f0f0e6d1ff3557": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "d929b9d20b1a4b00bb5902c0883b4034": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_e9113eab130140f9b9367a61b1e4b5c0",
-      "placeholder": "​",
-      "style": "IPY_MODEL_454dfc6558ae4a1b83621ea1352fc051",
-      "value": " 67015/67349 [00:08&lt;00:00, 8511.23 examples/s]"
-     }
-    },
-    "d9a622959c0d4c45b0b1932d7c85893d": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_758811c0d8284c76984fc2ac974d835a",
-      "placeholder": "​",
-      "style": "IPY_MODEL_4e06b19ebe63441ebb6a9b1cd57e2d21",
-      "value": "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-train.tfrecord*...:  77%"
-     }
-    },
-    "d9d4d67205d340a8acb919b7c80a842a": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "e5d9c734fb3e412aabc03e5790b12382": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "e70cb8396ab341c1bb0d60d2ebdd708a": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "e9113eab130140f9b9367a61b1e4b5c0": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "e9a881d7fcb84bafb6a063009bf4dd54": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_module_version": "1.5.0",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_77c0fa586d4d4dcd982513c827799adb",
-      "placeholder": "​",
-      "style": "IPY_MODEL_a70a3e92a37247eeaad91ef1e140aa54",
-      "value": " 1/1 [00:01&lt;00:00,  1.11s/ url]"
-     }
-    },
-    "efdf1618f9224961b840bf58fb763876": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "fac49204e83e49c9920c29a61382d9b7": {
-     "model_module": "@jupyter-widgets/base",
-     "model_module_version": "1.2.0",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": "20px"
-     }
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "N-uuvlp2_ybw"
+      },
+      "execution_count": 10,
+      "outputs": []
     }
-   }
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
+  ]
+}
\ No newline at end of file
diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb
deleted file mode 100644
index 71bfbb049d..0000000000
--- a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training_new.ipynb
+++ /dev/null
@@ -1,401 +0,0 @@
-{
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "gpuClass": "standard",
-    "accelerator": "GPU"
-  },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "source": [
-        "# keras-nlp installation"
-      ],
-      "metadata": {
-        "id": "FKTVkreu3MGG"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 1,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "NtcTiJqG20LV",
-        "outputId": "d03cd2ad-50ea-4553-8c91-255bcfed9065"
-      },
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
-          ]
-        }
-      ],
-      "source": [
-        "!pip install -q git+https://github.com/keras-team/keras-nlp.git tensorflow --upgrade"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "# Imports"
-      ],
-      "metadata": {
-        "id": "ofxRv3CK3Pdp"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "import keras_nlp\n",
-        "import tensorflow as tf\n",
-        "from tensorflow import keras\n",
-        "import tensorflow_datasets as tfds"
-      ],
-      "metadata": {
-        "id": "AHJ-nbcT3ATU"
-      },
-      "execution_count": 2,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "#Data"
-      ],
-      "metadata": {
-        "id": "r86tfpfB3T6_"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Load Dataset\n",
-        "\n",
-        "train_ds, valid_ds = tfds.load(\n",
-        "    \"glue/sst2\",\n",
-        "    split=[\"train\", \"validation\"],\n",
-        "    batch_size=16,\n",
-        ")\n",
-        "\n",
-        "def split_features(x):\n",
-        "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
-        "    # (features, label), where features is a tuple consisting of all\n",
-        "    # features.\n",
-        "    features = x[\"sentence\"]\n",
-        "    label = x[\"label\"]\n",
-        "    return (features, label)\n",
-        "\n",
-        "train_ds = train_ds.map(\n",
-        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
-        ").prefetch(tf.data.AUTOTUNE)\n",
-        "valid_ds = valid_ds.map(\n",
-        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
-        ").prefetch(tf.data.AUTOTUNE)"
-      ],
-      "metadata": {
-        "id": "8UDVK92k3VKQ"
-      },
-      "execution_count": 3,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Look first training set batch\n",
-        "# The format is (string_tensor, label_tensor)\n",
-        "train_ds.take(1).get_single_element()"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "w8cfKvOl3aiT",
-        "outputId": "25c58ee3-88b2-4c07-9973-6bc0b966c875"
-      },
-      "execution_count": 4,
-      "outputs": [
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
-              " array([b'for the uninitiated plays better on video with the sound ',\n",
-              "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
-              "        b'company once again dazzle and delight us ',\n",
-              "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
-              "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
-              "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
-              "        b'only masochistic moviegoers need apply . ',\n",
-              "        b'convince almost everyone that it was put on the screen , just for them ',\n",
-              "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
-              "        b'his supple understanding of the role ',\n",
-              "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
-              "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
-              "        b'of the startling intimacy ', b'draggy ',\n",
-              "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
-              "       dtype=object)>,\n",
-              " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
-            ]
-          },
-          "metadata": {},
-          "execution_count": 4
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "#Model"
-      ],
-      "metadata": {
-        "id": "CZTYpu0y4Y6A"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Create the Classifier Model \n",
-        "# For more details please look https://keras.io/guides/keras_nlp/getting_started/\n",
-        "\n",
-        "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
-        "    \"bert_tiny_en_uncased\",\n",
-        "    num_classes=2,\n",
-        "    dropout=0.1\n",
-        ")\n",
-        "\n",
-        "# Add loss function, optimizer and metrics\n",
-        "classifier.compile(\n",
-        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
-        "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
-        "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
-        "    jit_compile=True,\n",
-        ")\n",
-        "   \n",
-        "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need \n",
-        "# to do it if we just want to train it on a downstream task.\n",
-        "classifier.layers[-1].build([None, ])\n",
-        "classifier.summary()"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "1KxlC7gY32MH",
-        "outputId": "b02c8c90-9f37-40b8-843b-5fcc66819e80"
-      },
-      "execution_count": 5,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Model: \"bert_classifier\"\n",
-            "__________________________________________________________________________________________________\n",
-            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
-            "==================================================================================================\n",
-            " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
-            "                                                                                                  \n",
-            " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
-            "                                                                                                  \n",
-            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
-            "                                                                                                  \n",
-            " bert_backbone (BertBackbone)   {'sequence_output':  4385920     ['padding_mask[0][0]',           \n",
-            "                                 (None, None, 128),               'segment_ids[0][0]',            \n",
-            "                                 'pooled_output': (               'token_ids[0][0]']              \n",
-            "                                None, 128)}                                                       \n",
-            "                                                                                                  \n",
-            " dropout (Dropout)              (None, 128)          0           ['bert_backbone[0][0]']          \n",
-            "                                                                                                  \n",
-            " logits (Dense)                 (None, 2)            258         ['dropout[0][0]']                \n",
-            "                                                                                                  \n",
-            " bert_preprocessor (BertPreproc  multiple            0           []                               \n",
-            " essor)                                                                                           \n",
-            "                                                                                                  \n",
-            "==================================================================================================\n",
-            "Total params: 4,386,178\n",
-            "Trainable params: 4,386,178\n",
-            "Non-trainable params: 0\n",
-            "__________________________________________________________________________________________________\n"
-          ]
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Train the model\n",
-        "\n",
-        "N_EPOCHS = 2\n",
-        "classifier.fit(\n",
-        "    train_ds,\n",
-        "    validation_data=valid_ds,\n",
-        "    epochs=N_EPOCHS,\n",
-        ")\n",
-        "  "
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "fxVfW7t14wRS",
-        "outputId": "b9a4277d-4fa1-4668-9373-d8cef77bcc6a"
-      },
-      "execution_count": 6,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stderr",
-          "text": [
-            "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
-            "Instructions for updating:\n",
-            "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n"
-          ]
-        },
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Epoch 1/2\n",
-            "4210/4210 [==============================] - 828s 191ms/step - loss: 0.3782 - sparse_categorical_accuracy: 0.8299 - val_loss: 0.4344 - val_sparse_categorical_accuracy: 0.8165\n",
-            "Epoch 2/2\n",
-            "4210/4210 [==============================] - 783s 186ms/step - loss: 0.2409 - sparse_categorical_accuracy: 0.9039 - val_loss: 0.4626 - val_sparse_categorical_accuracy: 0.8222\n"
-          ]
-        },
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "<keras.callbacks.History at 0x7f2390120df0>"
-            ]
-          },
-          "metadata": {},
-          "execution_count": 6
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Save the weights\n",
-        "\n",
-        "model_name = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
-        "classifier.save_weights(f\"/content/{model_name}\")"
-      ],
-      "metadata": {
-        "id": "QExT3HYu6g6h"
-      },
-      "execution_count": 7,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Load the weigths\n",
-        "\n",
-        "weights_path = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
-        "classifier.load_weights(weights_path)"
-      ],
-      "metadata": {
-        "id": "8Y0ZTKRj62nq"
-      },
-      "execution_count": 8,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "#Test the model"
-      ],
-      "metadata": {
-        "id": "XWUIUhw47drP"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# We will shuffle the valid dataset and take 1st example\n",
-        "\n",
-        "shuffled_valid_ds = valid_ds.shuffle(55).rebatch(1)\n",
-        "element = shuffled_valid_ds.take(1).get_single_element()\n",
-        "\n",
-        "pred_logits = classifier.predict(element[0])\n",
-        "pred = tf.argmax(pred_logits)\n",
-        "\n",
-        "print(f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\")"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "FoQk4cy27jSk",
-        "outputId": "aff204c9-4879-403b-91ce-afcc3679b103"
-      },
-      "execution_count": 9,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "1/1 [==============================] - 3s 3s/step\n",
-            "Text :: preaches to two completely different choirs at the same time , which is a pretty amazing accomplishment .  \n",
-            "Label :: 1 \n",
-            "Model Prediction :: 1\n"
-          ]
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# You can test the model with your own statement!\n",
-        "\n",
-        "label_dict = {0:\"Bad Statement\", 1:\"Good Statement\"} # As for this dataset\n",
-        "\n",
-        "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
-        "output = tf.argmax(classifier.predict([\"Lord of the rings is best\"]), axis=1).numpy()[0]\n",
-        "print(label_dict[output])"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "7MB2U6Bh_j_m",
-        "outputId": "37a8223a-4178-4176-e660-258ed9b856c7"
-      },
-      "execution_count": 10,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "1/1 [==============================] - 1s 902ms/step\n",
-            "Good Statement\n"
-          ]
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [],
-      "metadata": {
-        "id": "N-uuvlp2_ybw"
-      },
-      "execution_count": 10,
-      "outputs": []
-    }
-  ]
-}
\ No newline at end of file

From 1cfd3069899e3b91143c00a91b818587ecb30a2b Mon Sep 17 00:00:00 2001
From: Susnato Dhar <susnatodhar10@gmail.com>
Date: Mon, 6 Mar 2023 10:20:29 +0530
Subject: [PATCH 3/5] nits

---
 .../bert_tiny_uncased_en_sst2_training.ipynb  | 750 +++++++++---------
 1 file changed, 370 insertions(+), 380 deletions(-)

diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
index 71bfbb049d..37d5c806c7 100644
--- a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
+++ b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
@@ -1,401 +1,391 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "gpuClass": "standard",
-    "accelerator": "GPU"
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "provenance": []
   },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "source": [
-        "# keras-nlp installation"
-      ],
-      "metadata": {
-        "id": "FKTVkreu3MGG"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 1,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "NtcTiJqG20LV",
-        "outputId": "d03cd2ad-50ea-4553-8c91-255bcfed9065"
-      },
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
-          ]
-        }
-      ],
-      "source": [
-        "!pip install -q git+https://github.com/keras-team/keras-nlp.git tensorflow --upgrade"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "# Imports"
-      ],
-      "metadata": {
-        "id": "ofxRv3CK3Pdp"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "import keras_nlp\n",
-        "import tensorflow as tf\n",
-        "from tensorflow import keras\n",
-        "import tensorflow_datasets as tfds"
-      ],
-      "metadata": {
-        "id": "AHJ-nbcT3ATU"
-      },
-      "execution_count": 2,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "#Data"
-      ],
-      "metadata": {
-        "id": "r86tfpfB3T6_"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Load Dataset\n",
-        "\n",
-        "train_ds, valid_ds = tfds.load(\n",
-        "    \"glue/sst2\",\n",
-        "    split=[\"train\", \"validation\"],\n",
-        "    batch_size=16,\n",
-        ")\n",
-        "\n",
-        "def split_features(x):\n",
-        "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
-        "    # (features, label), where features is a tuple consisting of all\n",
-        "    # features.\n",
-        "    features = x[\"sentence\"]\n",
-        "    label = x[\"label\"]\n",
-        "    return (features, label)\n",
-        "\n",
-        "train_ds = train_ds.map(\n",
-        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
-        ").prefetch(tf.data.AUTOTUNE)\n",
-        "valid_ds = valid_ds.map(\n",
-        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
-        ").prefetch(tf.data.AUTOTUNE)"
-      ],
-      "metadata": {
-        "id": "8UDVK92k3VKQ"
-      },
-      "execution_count": 3,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Look first training set batch\n",
-        "# The format is (string_tensor, label_tensor)\n",
-        "train_ds.take(1).get_single_element()"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "w8cfKvOl3aiT",
-        "outputId": "25c58ee3-88b2-4c07-9973-6bc0b966c875"
-      },
-      "execution_count": 4,
-      "outputs": [
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
-              " array([b'for the uninitiated plays better on video with the sound ',\n",
-              "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
-              "        b'company once again dazzle and delight us ',\n",
-              "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
-              "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
-              "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
-              "        b'only masochistic moviegoers need apply . ',\n",
-              "        b'convince almost everyone that it was put on the screen , just for them ',\n",
-              "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
-              "        b'his supple understanding of the role ',\n",
-              "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
-              "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
-              "        b'of the startling intimacy ', b'draggy ',\n",
-              "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
-              "       dtype=object)>,\n",
-              " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
-            ]
-          },
-          "metadata": {},
-          "execution_count": 4
-        }
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "#Model"
-      ],
-      "metadata": {
-        "id": "CZTYpu0y4Y6A"
-      }
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "gpuClass": "standard",
+  "accelerator": "GPU"
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# keras-nlp installation"
+   ],
+   "metadata": {
+    "id": "FKTVkreu3MGG"
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
+    "id": "NtcTiJqG20LV",
+    "outputId": "d03cd2ad-50ea-4553-8c91-255bcfed9065"
+   },
+   "outputs": [
     {
-      "cell_type": "code",
-      "source": [
-        "# Create the Classifier Model \n",
-        "# For more details please look https://keras.io/guides/keras_nlp/getting_started/\n",
-        "\n",
-        "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
-        "    \"bert_tiny_en_uncased\",\n",
-        "    num_classes=2,\n",
-        "    dropout=0.1\n",
-        ")\n",
-        "\n",
-        "# Add loss function, optimizer and metrics\n",
-        "classifier.compile(\n",
-        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
-        "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
-        "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
-        "    jit_compile=True,\n",
-        ")\n",
-        "   \n",
-        "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need \n",
-        "# to do it if we just want to train it on a downstream task.\n",
-        "classifier.layers[-1].build([None, ])\n",
-        "classifier.summary()"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "1KxlC7gY32MH",
-        "outputId": "b02c8c90-9f37-40b8-843b-5fcc66819e80"
-      },
-      "execution_count": 5,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Model: \"bert_classifier\"\n",
-            "__________________________________________________________________________________________________\n",
-            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
-            "==================================================================================================\n",
-            " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
-            "                                                                                                  \n",
-            " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
-            "                                                                                                  \n",
-            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
-            "                                                                                                  \n",
-            " bert_backbone (BertBackbone)   {'sequence_output':  4385920     ['padding_mask[0][0]',           \n",
-            "                                 (None, None, 128),               'segment_ids[0][0]',            \n",
-            "                                 'pooled_output': (               'token_ids[0][0]']              \n",
-            "                                None, 128)}                                                       \n",
-            "                                                                                                  \n",
-            " dropout (Dropout)              (None, 128)          0           ['bert_backbone[0][0]']          \n",
-            "                                                                                                  \n",
-            " logits (Dense)                 (None, 2)            258         ['dropout[0][0]']                \n",
-            "                                                                                                  \n",
-            " bert_preprocessor (BertPreproc  multiple            0           []                               \n",
-            " essor)                                                                                           \n",
-            "                                                                                                  \n",
-            "==================================================================================================\n",
-            "Total params: 4,386,178\n",
-            "Trainable params: 4,386,178\n",
-            "Non-trainable params: 0\n",
-            "__________________________________________________________________________________________________\n"
-          ]
-        }
-      ]
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "  Preparing metadata (setup.py) ... \u001B[?25l\u001B[?25hdone\n"
+     ]
+    }
+   ],
+   "source": [
+    "!pip install -q git+https://github.com/keras-team/keras-nlp.git tensorflow --upgrade"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Imports"
+   ],
+   "metadata": {
+    "id": "ofxRv3CK3Pdp"
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "import keras_nlp\n",
+    "import tensorflow as tf\n",
+    "from tensorflow import keras\n",
+    "import tensorflow_datasets as tfds"
+   ],
+   "metadata": {
+    "id": "AHJ-nbcT3ATU"
+   },
+   "execution_count": 2,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#Data"
+   ],
+   "metadata": {
+    "id": "r86tfpfB3T6_"
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# Load Dataset\n",
+    "\n",
+    "train_ds, valid_ds = tfds.load(\n",
+    "    \"glue/sst2\",\n",
+    "    split=[\"train\", \"validation\"],\n",
+    "    batch_size=16,\n",
+    ")\n",
+    "\n",
+    "def split_features(x):\n",
+    "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
+    "    # (features, label), where features is a tuple consisting of all\n",
+    "    # features.\n",
+    "    features = x[\"sentence\"]\n",
+    "    label = x[\"label\"]\n",
+    "    return (features, label)\n",
+    "\n",
+    "train_ds = train_ds.map(\n",
+    "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+    ").prefetch(tf.data.AUTOTUNE)\n",
+    "valid_ds = valid_ds.map(\n",
+    "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+    ").prefetch(tf.data.AUTOTUNE)"
+   ],
+   "metadata": {
+    "id": "8UDVK92k3VKQ"
+   },
+   "execution_count": 3,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# Look first training set batch\n",
+    "# The format is (string_tensor, label_tensor)\n",
+    "train_ds.take(1).get_single_element()"
+   ],
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
+    "id": "w8cfKvOl3aiT",
+    "outputId": "25c58ee3-88b2-4c07-9973-6bc0b966c875"
+   },
+   "execution_count": 4,
+   "outputs": [
     {
-      "cell_type": "code",
-      "source": [
-        "# Train the model\n",
-        "\n",
-        "N_EPOCHS = 2\n",
-        "classifier.fit(\n",
-        "    train_ds,\n",
-        "    validation_data=valid_ds,\n",
-        "    epochs=N_EPOCHS,\n",
-        ")\n",
-        "  "
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "fxVfW7t14wRS",
-        "outputId": "b9a4277d-4fa1-4668-9373-d8cef77bcc6a"
-      },
-      "execution_count": 6,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stderr",
-          "text": [
-            "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
-            "Instructions for updating:\n",
-            "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n"
-          ]
-        },
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Epoch 1/2\n",
-            "4210/4210 [==============================] - 828s 191ms/step - loss: 0.3782 - sparse_categorical_accuracy: 0.8299 - val_loss: 0.4344 - val_sparse_categorical_accuracy: 0.8165\n",
-            "Epoch 2/2\n",
-            "4210/4210 [==============================] - 783s 186ms/step - loss: 0.2409 - sparse_categorical_accuracy: 0.9039 - val_loss: 0.4626 - val_sparse_categorical_accuracy: 0.8222\n"
-          ]
-        },
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "<keras.callbacks.History at 0x7f2390120df0>"
-            ]
-          },
-          "metadata": {},
-          "execution_count": 6
-        }
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": [
+       "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
+       " array([b'for the uninitiated plays better on video with the sound ',\n",
+       "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
+       "        b'company once again dazzle and delight us ',\n",
+       "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
+       "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
+       "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
+       "        b'only masochistic moviegoers need apply . ',\n",
+       "        b'convince almost everyone that it was put on the screen , just for them ',\n",
+       "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
+       "        b'his supple understanding of the role ',\n",
+       "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
+       "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
+       "        b'of the startling intimacy ', b'draggy ',\n",
+       "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
+       "       dtype=object)>,\n",
+       " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
       ]
+     },
+     "metadata": {},
+     "execution_count": 4
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#Model"
+   ],
+   "metadata": {
+    "id": "CZTYpu0y4Y6A"
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# Create the Classifier Model \n",
+    "# For more details please look https://keras.io/guides/keras_nlp/getting_started/\n",
+    "\n",
+    "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
+    "    \"bert_tiny_en_uncased\",\n",
+    "    num_classes=2,\n",
+    "    dropout=0.1\n",
+    ")\n",
+    "\n",
+    "# Add loss function, optimizer and metrics\n",
+    "classifier.compile(\n",
+    "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
+    "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
+    "    jit_compile=True,\n",
+    ")\n",
+    "   \n",
+    "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need \n",
+    "# to do it if we just want to train it on a downstream task.\n",
+    "classifier.layers[-1].build([None, ])\n",
+    "classifier.summary()"
+   ],
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
+    "id": "1KxlC7gY32MH",
+    "outputId": "b02c8c90-9f37-40b8-843b-5fcc66819e80"
+   },
+   "execution_count": 5,
+   "outputs": [
     {
-      "cell_type": "code",
-      "source": [
-        "# Save the weights\n",
-        "\n",
-        "model_name = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
-        "classifier.save_weights(f\"/content/{model_name}\")"
-      ],
-      "metadata": {
-        "id": "QExT3HYu6g6h"
-      },
-      "execution_count": 7,
-      "outputs": []
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Model: \"bert_classifier\"\n",
+      "__________________________________________________________________________________________________\n",
+      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+      "==================================================================================================\n",
+      " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
+      "                                                                                                  \n",
+      " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
+      "                                                                                                  \n",
+      " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
+      "                                                                                                  \n",
+      " bert_backbone (BertBackbone)   {'sequence_output':  4385920     ['padding_mask[0][0]',           \n",
+      "                                 (None, None, 128),               'segment_ids[0][0]',            \n",
+      "                                 'pooled_output': (               'token_ids[0][0]']              \n",
+      "                                None, 128)}                                                       \n",
+      "                                                                                                  \n",
+      " dropout (Dropout)              (None, 128)          0           ['bert_backbone[0][0]']          \n",
+      "                                                                                                  \n",
+      " logits (Dense)                 (None, 2)            258         ['dropout[0][0]']                \n",
+      "                                                                                                  \n",
+      " bert_preprocessor (BertPreproc  multiple            0           []                               \n",
+      " essor)                                                                                           \n",
+      "                                                                                                  \n",
+      "==================================================================================================\n",
+      "Total params: 4,386,178\n",
+      "Trainable params: 4,386,178\n",
+      "Non-trainable params: 0\n",
+      "__________________________________________________________________________________________________\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# Train the model\n",
+    "\n",
+    "N_EPOCHS = 2\n",
+    "classifier.fit(\n",
+    "    train_ds,\n",
+    "    validation_data=valid_ds,\n",
+    "    epochs=N_EPOCHS,\n",
+    ")"
+   ],
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
+    "id": "fxVfW7t14wRS",
+    "outputId": "b9a4277d-4fa1-4668-9373-d8cef77bcc6a"
+   },
+   "execution_count": 6,
+   "outputs": [
     {
-      "cell_type": "code",
-      "source": [
-        "# Load the weigths\n",
-        "\n",
-        "weights_path = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
-        "classifier.load_weights(weights_path)"
-      ],
-      "metadata": {
-        "id": "8Y0ZTKRj62nq"
-      },
-      "execution_count": 8,
-      "outputs": []
+     "output_type": "stream",
+     "name": "stderr",
+     "text": [
+      "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
+      "Instructions for updating:\n",
+      "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n"
+     ]
     },
     {
-      "cell_type": "markdown",
-      "source": [
-        "#Test the model"
-      ],
-      "metadata": {
-        "id": "XWUIUhw47drP"
-      }
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Epoch 1/2\n",
+      "4210/4210 [==============================] - 828s 191ms/step - loss: 0.3782 - sparse_categorical_accuracy: 0.8299 - val_loss: 0.4344 - val_sparse_categorical_accuracy: 0.8165\n",
+      "Epoch 2/2\n",
+      "4210/4210 [==============================] - 783s 186ms/step - loss: 0.2409 - sparse_categorical_accuracy: 0.9039 - val_loss: 0.4626 - val_sparse_categorical_accuracy: 0.8222\n"
+     ]
     },
     {
-      "cell_type": "code",
-      "source": [
-        "# We will shuffle the valid dataset and take 1st example\n",
-        "\n",
-        "shuffled_valid_ds = valid_ds.shuffle(55).rebatch(1)\n",
-        "element = shuffled_valid_ds.take(1).get_single_element()\n",
-        "\n",
-        "pred_logits = classifier.predict(element[0])\n",
-        "pred = tf.argmax(pred_logits)\n",
-        "\n",
-        "print(f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\")"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "FoQk4cy27jSk",
-        "outputId": "aff204c9-4879-403b-91ce-afcc3679b103"
-      },
-      "execution_count": 9,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "1/1 [==============================] - 3s 3s/step\n",
-            "Text :: preaches to two completely different choirs at the same time , which is a pretty amazing accomplishment .  \n",
-            "Label :: 1 \n",
-            "Model Prediction :: 1\n"
-          ]
-        }
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": [
+       "<keras.callbacks.History at 0x7f2390120df0>"
       ]
+     },
+     "metadata": {},
+     "execution_count": 6
+    }
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# Save the weights\n",
+    "\n",
+    "model_name = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
+    "classifier.save_weights(f\"/content/{model_name}\")"
+   ],
+   "metadata": {
+    "id": "QExT3HYu6g6h"
+   },
+   "execution_count": 7,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# Load the weigths\n",
+    "\n",
+    "weights_path = f\"bert_tiny_uncased_en_sst2-epochs={N_EPOCHS}.h5\"\n",
+    "classifier.load_weights(weights_path)"
+   ],
+   "metadata": {
+    "id": "8Y0ZTKRj62nq"
+   },
+   "execution_count": 8,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#Test the model"
+   ],
+   "metadata": {
+    "id": "XWUIUhw47drP"
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# We will shuffle the valid dataset and take 1st example\n",
+    "\n",
+    "shuffled_valid_ds = valid_ds.shuffle(55).rebatch(1)\n",
+    "element = shuffled_valid_ds.take(1).get_single_element()\n",
+    "\n",
+    "pred_logits = classifier.predict(element[0])\n",
+    "pred = tf.argmax(pred_logits)\n",
+    "\n",
+    "print(f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\")"
+   ],
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
+    "id": "FoQk4cy27jSk",
+    "outputId": "aff204c9-4879-403b-91ce-afcc3679b103"
+   },
+   "execution_count": 9,
+   "outputs": [
     {
-      "cell_type": "code",
-      "source": [
-        "# You can test the model with your own statement!\n",
-        "\n",
-        "label_dict = {0:\"Bad Statement\", 1:\"Good Statement\"} # As for this dataset\n",
-        "\n",
-        "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
-        "output = tf.argmax(classifier.predict([\"Lord of the rings is best\"]), axis=1).numpy()[0]\n",
-        "print(label_dict[output])"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "7MB2U6Bh_j_m",
-        "outputId": "37a8223a-4178-4176-e660-258ed9b856c7"
-      },
-      "execution_count": 10,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "1/1 [==============================] - 1s 902ms/step\n",
-            "Good Statement\n"
-          ]
-        }
-      ]
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "1/1 [==============================] - 3s 3s/step\n",
+      "Text :: preaches to two completely different choirs at the same time , which is a pretty amazing accomplishment .  \n",
+      "Label :: 1 \n",
+      "Model Prediction :: 1\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# You can test the model with your own statement!\n",
+    "\n",
+    "label_dict = {0:\"Bad Statement\", 1:\"Good Statement\"} # As for this dataset\n",
+    "\n",
+    "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
+    "output = tf.argmax(classifier.predict([\"Lord of the rings is best\"]), axis=1).numpy()[0]\n",
+    "print(label_dict[output])"
+   ],
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
+    "id": "7MB2U6Bh_j_m",
+    "outputId": "37a8223a-4178-4176-e660-258ed9b856c7"
+   },
+   "execution_count": 10,
+   "outputs": [
     {
-      "cell_type": "code",
-      "source": [],
-      "metadata": {
-        "id": "N-uuvlp2_ybw"
-      },
-      "execution_count": 10,
-      "outputs": []
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "1/1 [==============================] - 1s 902ms/step\n",
+      "Good Statement\n"
+     ]
     }
-  ]
-}
\ No newline at end of file
+   ]
+  }
+ ]
+}

From 11b2cd7086fba03b9f80e5836ffaf6073f1862b9 Mon Sep 17 00:00:00 2001
From: Susnato Dhar <susnatodhar10@gmail.com>
Date: Tue, 7 Mar 2023 10:58:09 +0530
Subject: [PATCH 4/5] notebook check fixed

---
 .../bert_tiny_uncased_en_sst2_training.ipynb  | 28 ++++++++++++-------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
index 37d5c806c7..d331e8d048 100644
--- a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
+++ b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
@@ -91,6 +91,7 @@
     "    batch_size=16,\n",
     ")\n",
     "\n",
+    "\n",
     "def split_features(x):\n",
     "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
     "    # (features, label), where features is a tuple consisting of all\n",
@@ -99,6 +100,7 @@
     "    label = x[\"label\"]\n",
     "    return (features, label)\n",
     "\n",
+    "\n",
     "train_ds = train_ds.map(\n",
     "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
     ").prefetch(tf.data.AUTOTUNE)\n",
@@ -168,13 +170,11 @@
   {
    "cell_type": "code",
    "source": [
-    "# Create the Classifier Model \n",
+    "# Create the Classifier Model\n",
     "# For more details please look https://keras.io/guides/keras_nlp/getting_started/\n",
     "\n",
     "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
-    "    \"bert_tiny_en_uncased\",\n",
-    "    num_classes=2,\n",
-    "    dropout=0.1\n",
+    "    \"bert_tiny_en_uncased\", num_classes=2, dropout=0.1\n",
     ")\n",
     "\n",
     "# Add loss function, optimizer and metrics\n",
@@ -184,10 +184,14 @@
     "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
     "    jit_compile=True,\n",
     ")\n",
-    "   \n",
-    "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need \n",
+    "\n",
+    "# To see the summary(layers) of the model we need to build it(or call it once) but we don't need\n",
     "# to do it if we just want to train it on a downstream task.\n",
-    "classifier.layers[-1].build([None, ])\n",
+    "classifier.layers[-1].build(\n",
+    "    [\n",
+    "        None,\n",
+    "    ]\n",
+    ")\n",
     "classifier.summary()"
    ],
    "metadata": {
@@ -334,7 +338,9 @@
     "pred_logits = classifier.predict(element[0])\n",
     "pred = tf.argmax(pred_logits)\n",
     "\n",
-    "print(f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\")"
+    "print(\n",
+    "    f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\"\n",
+    ")"
    ],
    "metadata": {
     "colab": {
@@ -362,10 +368,12 @@
    "source": [
     "# You can test the model with your own statement!\n",
     "\n",
-    "label_dict = {0:\"Bad Statement\", 1:\"Good Statement\"} # As for this dataset\n",
+    "label_dict = {0: \"Bad Statement\", 1: \"Good Statement\"}  # As for this dataset\n",
     "\n",
     "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
-    "output = tf.argmax(classifier.predict([\"Lord of the rings is best\"]), axis=1).numpy()[0]\n",
+    "output = tf.argmax(\n",
+    "    classifier.predict([\"Lord of the rings is best\"]), axis=1\n",
+    ").numpy()[0]\n",
     "print(label_dict[output])"
    ],
    "metadata": {

From bceae22f612f769c57c99a1dfcfdafa906750eb1 Mon Sep 17 00:00:00 2001
From: Susnato Dhar <susnatodhar10@gmail.com>
Date: Tue, 7 Mar 2023 11:07:29 +0530
Subject: [PATCH 5/5] Small change

---
 .../bert_tiny_uncased_en_sst2_training.ipynb                  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
index d331e8d048..99277ab5cf 100644
--- a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
+++ b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
@@ -339,7 +339,8 @@
     "pred = tf.argmax(pred_logits)\n",
     "\n",
     "print(\n",
-    "    f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\"\n",
+    "    f\"Text :: {element[0].numpy()[0].decode('utf-8')} \\nLabel :: {element[1].numpy()[0]} \"\n",
+    "    f\"\\nModel Prediction :: {tf.argmax(pred_logits, axis=1).numpy()[0]}\"\n",
     ")"
    ],
    "metadata": {
@@ -370,7 +371,6 @@
     "\n",
     "label_dict = {0: \"Bad Statement\", 1: \"Good Statement\"}  # As for this dataset\n",
     "\n",
-    "text = \"The Lord of the rings is the Best, fantasy can offer!\"\n",
     "output = tf.argmax(\n",
     "    classifier.predict([\"Lord of the rings is best\"]), axis=1\n",
     ").numpy()[0]\n",