notebooks/cis6930_week9a_pretrained_lm_bert.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "cis6930-week9a-pretrained-lm-bert.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tydRDo6t2bjd"
      },
      "source": [
        "# CIS6930 Week 9a: Pre-trained Language Models (1) (Student version)\n",
        "\n",
        "---\n",
        "\n",
        "Preparation: Go to `Runtime > Change runtime type` and choose `GPU` for the hardware accelerator.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5MMZ5mxJ2iUy"
      },
      "source": [
        "gpu_info = !nvidia-smi -L\n",
        "gpu_info = \"\\n\".join(gpu_info)\n",
        "if gpu_info.find(\"failed\") >= 0:\n",
        "    print(\"Not connected to a GPU\")\n",
        "else:\n",
        "    print(gpu_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YbAny3QK2j6B"
      },
      "source": [
        "## Preparation\n",
        "\n",
        "For this notebookt, we use Hugging Face's `transformers` library."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1qK5hsVkyVF_"
      },
      "source": [
        "!pip install transformers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2sdU1QLtf00j"
      },
      "source": [
        "import copy\n",
        "from time import time\n",
        "from typing import Any, Dict\n",
        "import random\n",
        "\n",
        "import numpy as np\n",
        "from sklearn.metrics import accuracy_score\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch import optim\n",
        "from torch.utils.data import Dataset, TensorDataset, DataLoader\n",
        "from tqdm import tqdm\n",
        "from transformers import AdamW, get_linear_schedule_with_warmup"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWSMFA_g1PL9"
      },
      "source": [
        "## Playing with a Pre-trained Tokenizer\n",
        "\n",
        "As discussed in the lecture, we use \"pre-trained\" tokenizer models for pre-trained language models. Let's take a look at the tokenizer pre-trained for `bert-base-uncased`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZW_ATtDRe0C-"
      },
      "source": [
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
        "# AutoTokenizer, AutoModelForSequenceClassification are a \"meta\" class for \n",
        "# model-specific classes such as BertTokenizer, BertForSequenceClassification\n",
        "\n",
        "# See also Hugging Face's ModelHub\n",
        "# https://huggingface.co/bert-base-uncased\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z9pcd3MvuRWB"
      },
      "source": [
        "tokenizer.encode(\"Hello, world!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yDQI7RLW1tJK"
      },
      "source": [
        "for token in tokenizer.encode(\"Hello, world!\"):\n",
        "    print(tokenizer.decode(token))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tqmdVWu5uWla"
      },
      "source": [
        "tokenizer.decode(tokenizer.encode(\"Hello, world!\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VqxbCN0-uhjI"
      },
      "source": [
        "# \"Pre-trained\" vocabulary set\n",
        "len(tokenizer.vocab)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g3SysxuAvOyU"
      },
      "source": [
        "# Registered special tokens\n",
        "tokenizer.special_tokens_map"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gd4lVt7HuF54"
      },
      "source": [
        "# What about the longest English word? :)\n",
        "for tokenid in tokenizer.encode(\"pneumonoultramicroscopicsilicovolcanoconiosis\"):\n",
        "    print(tokenizer.decode(tokenid))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ryPSolUe2JRE"
      },
      "source": [
        "Now, we see the pre-trained tokenizer does not have the OoV issue and tokenzie an input sequence into subwords."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sde3-QPkwaBR"
      },
      "source": [
        "### `transformers.Tokenizer.__call__()`\n",
        "\n",
        "We usually use the `__call__()` method. The method returns a dictionary of `input_ids`, `token_type_ids`, and `attention_mask`, which are compatible with the interface of pre-trained language models in the `transformers` library. This function is also convenient for **padding**.\n",
        "\n",
        "Let's take a look."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mo0rtzDUveuv"
      },
      "source": [
        "tokenizer(\"Hello, world!\")  # __call__() in Python"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r8waI4QBvkIV"
      },
      "source": [
        "# For sequence-pair classification\n",
        "tokenizer(\"Hello, world!\", \"Good morning world!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hzTLoFPXxE5O"
      },
      "source": [
        "# By providing `max_length` and `padding` arguments,\n",
        "# you can make \"pre-padded\" token ID sequences\n",
        "tokenizer([\"Hello, world!\", \"Hello, again!\"],\n",
        "          max_length=16, padding=\"max_length\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wZgFZtcOxdUH"
      },
      "source": [
        "# By adding `return_tensors=\"pt\"`, you can get PyTorch tensor objects instead of lists. \n",
        "tokenizer([\"Hello, world!\", \"Hello, again!\"],\n",
        "          max_length=16, padding=\"max_length\",\n",
        "          return_tensors=\"pt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5jB4eGHyfEl"
      },
      "source": [
        "## Implementing Custom Dataset class\n",
        "\n",
        "In Week 6, we created a custom dataset for the Twitter dataset (Please see [the Google Colab notebook](https://colab.research.google.com/drive/1DZN-Bo2HBnPQPm4jrQzEIchhHdN682qP?usp=sharing))\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O4WGTpfRez-s"
      },
      "source": [
        "# https://www.kaggle.com/crowdflower/twitter-airline-sentiment\n",
        "# License CC BY-NC-SA 4.0\n",
        "!gdown --id 1BS_TIqm7crkBRr8p6REZrMv4Uk9_-e6W"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jnxo0_CBe0Aw"
      },
      "source": [
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "\n",
        "from torch.utils.data import Dataset, TensorDataset, DataLoader\n",
        "\n",
        "# Loading dataset\n",
        "df = pd.read_csv(\"Tweets.csv\")\n",
        "\n",
        "# Label encoder\n",
        "le = LabelEncoder()\n",
        "y = le.fit_transform(df[\"airline_sentiment\"].values)\n",
        "df[\"label\"] = y\n",
        "\n",
        "# Splint into 60% train, 20% valid, 20% test\n",
        "train_df, test_df = train_test_split(\n",
        "    df, test_size=0.2, random_state=1)\n",
        "\n",
        "train_df, valid_df = train_test_split(\n",
        "    train_df, test_size=0.25, random_state=1)  # 0.25 x 0.8 = 0.2\n",
        "\n",
        "print(len(train_df), len(valid_df), len(test_df))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yc77syuIXcoe"
      },
      "source": [
        "df.columns"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7aJuP3eXe0FS"
      },
      "source": [
        "class TweetDataset(Dataset):\n",
        "    def __init__(self,\n",
        "                 df,\n",
        "                 tokenizer,\n",
        "                 max_length=256):\n",
        "        self.df = df\n",
        "        input_ids = []\n",
        "        for text in self.df[\"text\"].tolist():\n",
        "            d = tokenizer(text,\n",
        "                          max_length=max_length,\n",
        "                          padding=\"max_length\",\n",
        "                          return_tensors=\"pt\")\n",
        "            for k, v in d.items():\n",
        "                # To remove unnecessary list\n",
        "                d[k] = v.squeeze(0)\n",
        "            input_ids.append(d)\n",
        "\n",
        "        self.df[\"input_ids\"] = input_ids\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.df)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        return {**self.df.iloc[idx][\"input_ids\"],\n",
        "                \"labels\": self.df.iloc[idx][\"label\"]}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CRShdaokI5lP"
      },
      "source": [
        "train_dataset = TweetDataset(train_df, tokenizer, max_length=256)\n",
        "valid_dataset = TweetDataset(valid_df, tokenizer, max_length=256)\n",
        "test_dataset = TweetDataset(test_df, tokenizer, max_length=256)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yj3EGXavJMyf"
      },
      "source": [
        "# Take a look at a sample batch\n",
        "batch = next(iter(DataLoader(train_dataset, batch_size=4)))\n",
        "batch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QhbirlG12mrH"
      },
      "source": [
        "### Testing a Pre-trained Model (with a sequence classification head)\n",
        "\n",
        "- `AutoModel`: Base model (Note that it does not have any classification heads)\n",
        "- `AutoModelForSequenceClasiffication`: For single-sentence/sentnece-pair classification \n",
        "- `AutoModelForTokenClassification`: For sequential tagging\n",
        "- `AutoModelForQuestionAnswering`: For Question Answering\n",
        "\n",
        "See [the official API documentation](https://huggingface.co/transformers/model_doc/auto.html) for details. \n",
        "\n",
        "In this example, we will use `AutoModelForSequenceClassification` for a text classification problem.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "in-kZVFfNooF"
      },
      "source": [
        "model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PiEHB8ab3rgw"
      },
      "source": [
        "Let's take a look at how the `forward()` function is implemented for `AutoModelForSequentialClassification`, which is actually `BertModelForSequenceClassification` in this case (as we load a pre-trained BERTmodel) \n",
        "\n",
        "https://huggingface.co/transformers/model_doc/bert.html#transformers.BertModel.forward\n",
        "\n",
        "When `forward()` takes the optional argument `labels`, it will return the corresponding `loss` value. This is convenient and now we don't have to manually calculate the loss value. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qiwrC1HUPXrG"
      },
      "source": [
        "output = model(**batch)\n",
        "output"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wgQf3hC2Z0BM"
      },
      "source": [
        "output.loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WhJtkyye4yY2"
      },
      "source": [
        "## Training script\n",
        "\n",
        "The following training script is based on the previous version with little modifiations:\n",
        "- 1) Replaced manual loss calculation with model's output.\n",
        "- 2) Added a learning rate scheduler.\n",
        "\n",
        "Let's take a look."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qpid7CUjS38Q"
      },
      "source": [
        "def train(model: nn.Module,\n",
        "          train_dataset: Dataset,\n",
        "          valid_dataset: Dataset,\n",
        "          config: Dict[str, Any],\n",
        "          random_seed: int = 0):\n",
        "  \n",
        "    # Random Seeds ===============\n",
        "    torch.manual_seed(random_seed)\n",
        "    random.seed(random_seed)\n",
        "    np.random.seed(random_seed)\n",
        "    # Random Seeds ===============\n",
        "\n",
        "    # GPU configuration\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "    dl_train = DataLoader(train_dataset,\n",
        "                          batch_size=config[\"batch_size\"],\n",
        "                          shuffle=True,\n",
        "                          drop_last=True)\n",
        "    dl_valid = DataLoader(valid_dataset)\n",
        "                  \n",
        "    # Model, Optimzier, Loss function\n",
        "    model = model.to(device)\n",
        "\n",
        "    # Optimizer    \n",
        "    no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
        "    optimizer_grouped_parameters = [{\n",
        "            \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
        "            \"weight_decay\": 0.0\n",
        "        },\n",
        "        {\n",
        "            \"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],\n",
        "            \"weight_decay\": 0.0},\n",
        "        ]\n",
        "    optimizer = config[\"optimizer_cls\"](optimizer_grouped_parameters,\n",
        "                                        lr=config[\"lr\"])\n",
        "    t_total = len(dl_train) * config[\"n_epochs\"]\n",
        "    scheduler = config[\"scheduler_cls\"](optimizer,\n",
        "                                        num_warmup_steps=0,\n",
        "                                        num_training_steps=t_total)\n",
        "    \n",
        "    # For each epoch\n",
        "    eval_list = []\n",
        "    t0 = time()\n",
        "    best_val = None\n",
        "    best_model = None\n",
        "    for n in range(config[\"n_epochs\"]):\n",
        "        t1 = time()\n",
        "        print(\"Epoch {}\".format(n))\n",
        "        # Training\n",
        "        train_loss = 0.\n",
        "        train_pred_list = []\n",
        "        train_true_list = []\n",
        "        model.train()  # Switch to the training mode\n",
        "\n",
        "        # For each batch\n",
        "        for batch in tqdm(dl_train):\n",
        "            optimizer.zero_grad()              # Initialize gradient information\n",
        "            # ==================================================================\n",
        "            for k, v in batch.items():\n",
        "                batch[k] = v.to(device)\n",
        "            output = model(**batch)\n",
        "            loss = output.loss\n",
        "            preds = output.logits.argmax(axis=1).detach().cpu().tolist()\n",
        "            labels = batch[\"labels\"].detach().cpu().tolist()\n",
        "            # ==================================================================            \n",
        "            loss.backward()                    # Backpropagate the loss value\n",
        "            optimizer.step()                   # Update the parameters\n",
        "            scheduler.step()                   # [New] Update the scheduler step\n",
        "            train_loss += loss.data.item()\n",
        "            train_pred_list += preds\n",
        "            train_true_list += labels\n",
        "\n",
        "        train_loss /= len(dl_train)\n",
        "        train_acc = accuracy_score(train_true_list, train_pred_list)\n",
        "        print(\"    Training loss: {:.4f}    Training acc: {:.4f}\".format(train_loss,\n",
        "                                                                         train_acc))\n",
        "\n",
        "        # Validation\n",
        "        valid_loss = 0.\n",
        "        valid_pred_list = []\n",
        "        valid_true_list = []\n",
        "\n",
        "        model.eval()  # Switch to the evaluation mode\n",
        "        valid_emb_list = []\n",
        "        valid_label_list = []\n",
        "        for i, batch in tqdm(enumerate(dl_valid)):\n",
        "            # ==================================================================\n",
        "            for k, v in batch.items():\n",
        "                batch[k] = v.to(device)\n",
        "            output = model(**batch)\n",
        "            loss = output.loss\n",
        "            preds = output.logits.argmax(axis=1).detach().cpu().tolist()\n",
        "            labels = batch[\"labels\"].detach().cpu().tolist()\n",
        "            # ==================================================================\n",
        "            valid_loss += loss.data.item()\n",
        "            valid_pred_list += preds\n",
        "            valid_true_list += labels\n",
        "\n",
        "        valid_loss /= len(dl_valid)\n",
        "        valid_acc = accuracy_score(valid_true_list, valid_pred_list)\n",
        "        print(\"  Validation loss: {:.4f}  Validation acc: {:.4f}\".format(valid_loss,\n",
        "                                                                         valid_acc))\n",
        "\n",
        "        # Model selection\n",
        "        if best_val is None or valid_loss < best_val:\n",
        "            best_model = copy.deepcopy(model)\n",
        "            best_val = valid_loss\n",
        "\n",
        "        t2 = time()\n",
        "        print(\"     Elapsed time: {:.1f} [sec]\".format(t2 - t1))\n",
        "\n",
        "        # Store train/validation loss values\n",
        "        eval_list.append([n, train_loss, valid_loss, train_acc, valid_acc, t2 - t1])\n",
        "\n",
        "    eval_df = pd.DataFrame(eval_list, columns=[\"epoch\",\n",
        "                                               \"train_loss\", \"valid_loss\",\n",
        "                                               \"train_acc\", \"valid_acc\",\n",
        "                                               \"time\"])\n",
        "    eval_df.set_index(\"epoch\")\n",
        "\n",
        "    print(\"Total time: {:.1f} [sec]\".format(t2 - t0))\n",
        "\n",
        "    # Return the best model and trainining/validation information\n",
        "    return {\"model\": best_model,\n",
        "            \"best_val\": best_val,\n",
        "            \"eval_df\": eval_df}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q1fftFjs5fIA"
      },
      "source": [
        "## Let's try!\n",
        "\n",
        "Let's run the trainin script with the following configuration. Note that it several minutes to finish one epoch. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tnlbWoj1au28"
      },
      "source": [
        "config = {\"optimizer_cls\": optim.AdamW,\n",
        "          \"scheduler_cls\": get_linear_schedule_with_warmup,\n",
        "          \"lr\": 5e-5,\n",
        "          \"batch_size\": 8,\n",
        "          \"n_epochs\": 3}\n",
        "model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=3)\n",
        "output = train(model, train_dataset, valid_dataset, config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pnJUZROF8Rte"
      },
      "source": [
        "### Results\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k26Apt_-n_dJ"
      },
      "source": [
        "output[\"eval_df\"]"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}