5-bert_transfer_learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT Encoder + Classifier Head Prediction\n",
    "\n",
    "In this final notebook, we add a custom classification neural network head onto the pre-trained BERT encoder and train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "[]\n"
     ]
    }
   ],
   "source": [
    "# General Imports\n",
    "import pandas as pd \n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.utils import shuffle\n",
    "from os.path import join\n",
    "\n",
    "# NN-related imports\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub \n",
    "import tensorflow_text as text \n",
    "\n",
    "print(tf.test.is_built_with_cuda())\n",
    "print(tf.config.list_physical_devices('GPU'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_preprocess = hub.KerasLayer('https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3')\n",
    "bert = hub.KerasLayer('https://tfhub.dev/google/experts/bert/wiki_books/sst2/2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Rating</th>\n",
       "      <th>Review</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>This album is a travesty to the songs of the 5...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>I found this book a complete waist of time, I ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>The product I got had scratches on its surface...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>Ok well it may not be the worst book that I ha...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>It was ok. Could of gotten into the other char...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49995</th>\n",
       "      <td>5</td>\n",
       "      <td>Good little memory stick. Currently using as m...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49996</th>\n",
       "      <td>5</td>\n",
       "      <td>How anyone can write such fun tropical songs a...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49997</th>\n",
       "      <td>5</td>\n",
       "      <td>Sure, its one of The Great Man's best movies, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49998</th>\n",
       "      <td>5</td>\n",
       "      <td>I finally bought and watched this classic epic...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49999</th>\n",
       "      <td>5</td>\n",
       "      <td>I'm a little late in posting a review, one whi...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>50000 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       Rating                                             Review\n",
       "0           1  This album is a travesty to the songs of the 5...\n",
       "1           1  I found this book a complete waist of time, I ...\n",
       "2           1  The product I got had scratches on its surface...\n",
       "3           1  Ok well it may not be the worst book that I ha...\n",
       "4           1  It was ok. Could of gotten into the other char...\n",
       "...       ...                                                ...\n",
       "49995       5  Good little memory stick. Currently using as m...\n",
       "49996       5  How anyone can write such fun tropical songs a...\n",
       "49997       5  Sure, its one of The Great Man's best movies, ...\n",
       "49998       5  I finally bought and watched this classic epic...\n",
       "49999       5  I'm a little late in posting a review, one whi...\n",
       "\n",
       "[50000 rows x 2 columns]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_dir = \"data/\"\n",
    "data = pd.read_csv(join(data_dir, \"downsampled_train_50000.csv\"))[[\"Rating\", \"Review\"]]\n",
    "data[\"Review\"] = data[\"Review\"].apply(str)\n",
    "display(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model():\n",
    "    \"\"\"Build model with custom classifier stacked on top of BERT encoder\"\"\"\n",
    "    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='txt_input')\n",
    "    bert_input = bert_preprocess(text_input)\n",
    "    bert_output = bert(bert_input)\n",
    "    clf_input = bert_output['pooled_output']\n",
    "    clf = tf.keras.layers.Dropout(0.1)(clf_input) \n",
    "    clf = tf.keras.layers.Dense(384, activation='sigmoid', kernel_regularizer='l2')(clf)\n",
    "    clf = tf.keras.layers.Dropout(0.1)(clf)\n",
    "    clf = tf.keras.layers.Dense(5, activation='sigmoid', name='clf')(clf)\n",
    "    return tf.keras.Model(text_input, clf)\n",
    "\n",
    "model = build_model() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "txt_input (InputLayer)          [(None,)]            0                                            \n",
      "__________________________________________________________________________________________________\n",
      "keras_layer (KerasLayer)        {'input_mask': (None 0           txt_input[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "keras_layer_1 (KerasLayer)      {'sequence_output':  109482241   keras_layer[0][0]                \n",
      "                                                                 keras_layer[0][1]                \n",
      "                                                                 keras_layer[0][2]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout (Dropout)               (None, 768)          0           keras_layer_1[0][13]             \n",
      "__________________________________________________________________________________________________\n",
      "dense (Dense)                   (None, 384)          295296      dropout[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "dropout_1 (Dropout)             (None, 384)          0           dense[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "clf (Dense)                     (None, 5)            1925        dropout_1[0][0]                  \n",
      "==================================================================================================\n",
      "Total params: 109,779,462\n",
      "Trainable params: 297,221\n",
      "Non-trainable params: 109,482,241\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "metrics = [tf.metrics.SparseCategoricalAccuracy()]\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "model.compile(loss=loss, optimizer=optimizer, metrics=metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X_train: (7200,) | X_val: (1800,) | X_test: (1000,) | \n",
      "y_train: (7200,) | y_val: (1800,) | y_test: (1000,) | \n"
     ]
    }
   ],
   "source": [
    "train_data = shuffle(data)[:10000]\n",
    "X = train_data[\"Review\"].to_numpy()\n",
    "y = train_data[\"Rating\"].to_numpy() - 1\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1)\n",
    "X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2)\n",
    "print(f\"X_train: {X_train.shape} | X_val: {X_val.shape} | X_test: {X_test.shape} | \\n\" +\n",
    "    f\"y_train: {y_train.shape} | y_val: {y_val.shape} | y_test: {y_test.shape} | \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/25\n",
      "113/113 [==============================] - 1095s 10s/step - loss: 1.1777 - sparse_categorical_accuracy: 0.4704 - val_loss: 1.1010 - val_sparse_categorical_accuracy: 0.5028\n",
      "Epoch 2/25\n",
      "113/113 [==============================] - 1187s 11s/step - loss: 1.1136 - sparse_categorical_accuracy: 0.5046 - val_loss: 1.1129 - val_sparse_categorical_accuracy: 0.5017\n",
      "Epoch 3/25\n",
      "113/113 [==============================] - 1062s 9s/step - loss: 1.0880 - sparse_categorical_accuracy: 0.5189 - val_loss: 1.0774 - val_sparse_categorical_accuracy: 0.5294\n",
      "Epoch 4/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 1.0756 - sparse_categorical_accuracy: 0.5269 - val_loss: 1.0730 - val_sparse_categorical_accuracy: 0.5244\n",
      "Epoch 5/25\n",
      "113/113 [==============================] - 1059s 9s/step - loss: 1.0721 - sparse_categorical_accuracy: 0.5254 - val_loss: 1.0970 - val_sparse_categorical_accuracy: 0.5172\n",
      "Epoch 6/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 1.0538 - sparse_categorical_accuracy: 0.5418 - val_loss: 1.0746 - val_sparse_categorical_accuracy: 0.5183\n",
      "Epoch 7/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 1.0472 - sparse_categorical_accuracy: 0.5383 - val_loss: 1.0966 - val_sparse_categorical_accuracy: 0.5122\n",
      "Epoch 8/25\n",
      "113/113 [==============================] - 1061s 9s/step - loss: 1.0441 - sparse_categorical_accuracy: 0.5421 - val_loss: 1.0653 - val_sparse_categorical_accuracy: 0.5344\n",
      "Epoch 9/25\n",
      "113/113 [==============================] - 1058s 9s/step - loss: 1.0375 - sparse_categorical_accuracy: 0.5503 - val_loss: 1.1156 - val_sparse_categorical_accuracy: 0.5078\n",
      "Epoch 10/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 1.0217 - sparse_categorical_accuracy: 0.5547 - val_loss: 1.0860 - val_sparse_categorical_accuracy: 0.5244\n",
      "Epoch 11/25\n",
      "113/113 [==============================] - 1059s 9s/step - loss: 1.0236 - sparse_categorical_accuracy: 0.5522 - val_loss: 1.0754 - val_sparse_categorical_accuracy: 0.5350\n",
      "Epoch 12/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 1.0096 - sparse_categorical_accuracy: 0.5593 - val_loss: 1.0875 - val_sparse_categorical_accuracy: 0.5200\n",
      "Epoch 13/25\n",
      "113/113 [==============================] - 1057s 9s/step - loss: 1.0083 - sparse_categorical_accuracy: 0.5665 - val_loss: 1.0920 - val_sparse_categorical_accuracy: 0.5156\n",
      "Epoch 14/25\n",
      "113/113 [==============================] - 1058s 9s/step - loss: 0.9940 - sparse_categorical_accuracy: 0.5718 - val_loss: 1.0719 - val_sparse_categorical_accuracy: 0.5317\n",
      "Epoch 15/25\n",
      "113/113 [==============================] - 1057s 9s/step - loss: 0.9894 - sparse_categorical_accuracy: 0.5669 - val_loss: 1.0796 - val_sparse_categorical_accuracy: 0.5283\n",
      "Epoch 16/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 0.9801 - sparse_categorical_accuracy: 0.5772 - val_loss: 1.0754 - val_sparse_categorical_accuracy: 0.5350\n",
      "Epoch 17/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 0.9736 - sparse_categorical_accuracy: 0.5782 - val_loss: 1.1032 - val_sparse_categorical_accuracy: 0.5128\n",
      "Epoch 18/25\n",
      "113/113 [==============================] - 1062s 9s/step - loss: 0.9727 - sparse_categorical_accuracy: 0.5775 - val_loss: 1.0899 - val_sparse_categorical_accuracy: 0.5156\n",
      "Epoch 19/25\n",
      "113/113 [==============================] - 1061s 9s/step - loss: 0.9596 - sparse_categorical_accuracy: 0.5929 - val_loss: 1.0951 - val_sparse_categorical_accuracy: 0.5183\n",
      "Epoch 20/25\n",
      "113/113 [==============================] - 1059s 9s/step - loss: 0.9505 - sparse_categorical_accuracy: 0.5907 - val_loss: 1.0939 - val_sparse_categorical_accuracy: 0.5128\n",
      "Epoch 21/25\n",
      "113/113 [==============================] - 1061s 9s/step - loss: 0.9466 - sparse_categorical_accuracy: 0.5914 - val_loss: 1.0978 - val_sparse_categorical_accuracy: 0.5133\n",
      "Epoch 22/25\n",
      "113/113 [==============================] - 1060s 9s/step - loss: 0.9351 - sparse_categorical_accuracy: 0.6017 - val_loss: 1.1037 - val_sparse_categorical_accuracy: 0.5128\n",
      "Epoch 23/25\n",
      "113/113 [==============================] - 1061s 9s/step - loss: 0.9282 - sparse_categorical_accuracy: 0.5975 - val_loss: 1.1057 - val_sparse_categorical_accuracy: 0.5133\n",
      "Epoch 24/25\n",
      "113/113 [==============================] - 1061s 9s/step - loss: 0.9086 - sparse_categorical_accuracy: 0.6118 - val_loss: 1.1387 - val_sparse_categorical_accuracy: 0.5028\n",
      "Epoch 25/25\n",
      "113/113 [==============================] - 1066s 9s/step - loss: 0.9008 - sparse_categorical_accuracy: 0.6179 - val_loss: 1.1067 - val_sparse_categorical_accuracy: 0.5133\n"
     ]
    }
   ],
   "source": [
    "EPOCHS = 25\n",
    "BATCH_SIZE = 64\n",
    "history = model.fit(X_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_data=(X_val, y_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 900). These functions will not be directly callable after loading.\n",
      "WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 900). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: models/bert2/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: models/bert2/assets\n"
     ]
    }
   ],
   "source": [
    "tf.saved_model.save(model, \"models/bert2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:5 out of the last 5 calls to <function recreate_function.<locals>.restored_function_body at 0x7fbfbda8d170> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:5 out of the last 5 calls to <function recreate_function.<locals>.restored_function_body at 0x7fbfbda8d170> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:6 out of the last 6 calls to <function recreate_function.<locals>.restored_function_body at 0x7fbfbda9ed40> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:6 out of the last 6 calls to <function recreate_function.<locals>.restored_function_body at 0x7fbfbda9ed40> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:7 out of the last 7 calls to <function recreate_function.<locals>.restored_function_body at 0x7fbfc61a4050> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:7 out of the last 7 calls to <function recreate_function.<locals>.restored_function_body at 0x7fbfc61a4050> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    }
   ],
   "source": [
    "model = tf.keras.models.load_model(\"models/bert2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32/32 [==============================] - 110s 3s/step - loss: 1.0776 - sparse_categorical_accuracy: 0.5340\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[1.077552318572998, 0.5339999794960022]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.evaluate(X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see this approach achieves a 53% accuracy on the test set, more than doubling the accuracy achieved by th baseline Logistic Regression model. The model's accuracy is still 12% less than the state-of-the-art approach which achieved a 65% accuracy (https://paperswithcode.com/sota/sentiment-analysis-on-amazon-review-full)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 64-bit ('askoski': conda)",
   "language": "python",
   "name": "python37564bitaskoskicondaf89acb10eb5d40d0a2b1becabd7f59e8"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  },
  "metadata": {
   "interpreter": {
    "hash": "7a72ecedb2cd54256b44a9751ae0cfff8f7b1c55503874d32066fee2d69e1c93"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}