examples/6_text_classification_task_using_bert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text Classification Task using BERT (ClickBait detection task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we will demonstrate how to fine-tune and evaluate the BERT model on the ClickBait dataset using PyTorchWrapper."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Downloading Data\n",
    "First of all we download and extract the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p data/\n",
    "\n",
    "! wget -P data/ https://github.com/bhargaviparanjape/clickbait/raw/master/dataset/clickbait_data.gz\n",
    "! wget -P data/ https://github.com/bhargaviparanjape/clickbait/raw/master/dataset/non_clickbait_data.gz\n",
    "\n",
    "! gunzip -f data/clickbait_data.gz \n",
    "! gunzip -f data/non_clickbait_data.gz \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Additional libraries\n",
    "\n",
    "Next we need to install the `tranformers` library in order use the pretrained BERT model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import Statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "import random\n",
    "import math\n",
    "\n",
    "from torch import nn\n",
    "from collections import Counter\n",
    "from glob import glob\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from torch.utils.data.sampler import SequentialSampler, SubsetRandomSampler\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import BertModel, BertTokenizer\n",
    "\n",
    "from pytorch_wrapper import modules, System\n",
    "from pytorch_wrapper import functional as pwF\n",
    "from pytorch_wrapper import evaluators as evaluators\n",
    "from pytorch_wrapper.loss_wrappers import GenericPointWiseLossWrapper\n",
    "from pytorch_wrapper.training_callbacks import EarlyStoppingCriterionCallback\n",
    "from pytorch_wrapper.samplers import SubsetOrderedBatchWiseRandomSampler, SubsetOrderedSequentialSampler, \\\n",
    "    OrderedSequentialSampler\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dataset Definition\n",
    "Next we create the ClickBaitDataset class. We will use the pretrained BPE tokenizer provided by the `transformers` library in order to prepare the input for the BERT model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "\n",
    "\n",
    "class ClickBaitDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        self.ids = []\n",
    "        self.texts = []\n",
    "        self.texts_len = []\n",
    "        self.targets = []\n",
    "\n",
    "        current_id = 0\n",
    "        for filename in ['clickbait_data', 'non_clickbait_data']:\n",
    "            with open(f'data/{filename}') as fw:\n",
    "                for line in tqdm(fw):\n",
    "                    if line == '\\n':\n",
    "                        continue\n",
    "                    self.ids.append(current_id)\n",
    "                    text = bert_tokenizer.encode(line.lower(), add_special_tokens=True)\n",
    "                    self.texts.append(text)\n",
    "                    self.texts_len.append(len(text))\n",
    "                    self.targets.append(filename == 'clickbait_data')\n",
    "\n",
    "        self._shuffle_examples()\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "\n",
    "        return (\n",
    "            self.ids[index],\n",
    "            (\n",
    "                self.texts[index],\n",
    "                self.texts_len[index]\n",
    "            ),\n",
    "            self.targets[index]\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ids)\n",
    "\n",
    "    def _shuffle_examples(self, seed=12345):\n",
    "        \"\"\"\n",
    "        Shuffles the examples with the given seed.\n",
    "        :param seed: The seed used for shuffling.\n",
    "        \"\"\"\n",
    "        random.seed(seed)\n",
    "        l = list(zip(self.ids, self.texts, self.texts_len, self.targets))\n",
    "        random.shuffle(l)\n",
    "        self.ids, self.texts, self.texts_len, self.targets = zip(*l)\n",
    "\n",
    "    @staticmethod\n",
    "    def collate_fn(batch):\n",
    "        \"\"\"\n",
    "        Function that combines a list of examples into a batch (Called internally by dataloaders).\n",
    "        \"\"\"\n",
    "        batch_zipped = list(zip(*batch))\n",
    "        input_zipped = list(zip(*batch_zipped[1]))\n",
    "\n",
    "        ids = batch_zipped[0]\n",
    "        texts = torch.tensor(ClickBaitDataset.pad_to_max(input_zipped[0]), dtype=torch.long)\n",
    "        texts_len = torch.tensor(input_zipped[1], dtype=torch.int)\n",
    "        targets = torch.tensor(batch_zipped[2], dtype=torch.float)\n",
    "\n",
    "        return {\n",
    "\n",
    "            'id': ids,\n",
    "            'input': [texts, texts_len],\n",
    "            'target': targets\n",
    "        }\n",
    "\n",
    "    @staticmethod\n",
    "    def pad_to_max(lst, max_len=None, pad_int=0):\n",
    "        \"\"\"\n",
    "        Pads the given list of list of tokens to the maximum length.\n",
    "        :param lst: List of list of tokens.\n",
    "        \"\"\"\n",
    "        pad = len(max(lst, key=len))\n",
    "        if max_len is not None:\n",
    "            pad = min(max_len, pad)\n",
    "\n",
    "        return [i + [pad_int] * (pad - len(i)) if len(i) <= pad else i[:pad] for i in lst]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model Definition\n",
    "In this example we will use the pretrained base uncased BERT model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BERTModel, self).__init__()\n",
    "\n",
    "        self.bert = BertModel.from_pretrained('bert-base-uncased')\n",
    "        self.output_linear = nn.Linear(768, 1)\n",
    "\n",
    "    def forward(self, text, text_len):\n",
    "        bert_last_hidden_states = self.bert(text)[0]\n",
    "        mask = pwF.create_mask_from_length(text_len, bert_last_hidden_states.shape[-2], zeros_at_end=True)\n",
    "        encoding = pwF.masked_mean_pooling(bert_last_hidden_states, mask, -2)\n",
    "        return self.output_linear(encoding).squeeze(-1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training\n",
    "\n",
    "Next we create the dataset object along with three data loaders (for training, validation, and testing). We will also make use of `SubsetOrderedBatchWiseRandomSampler` and `SubsetOrderedSequentialSampler` in order to batch together texts with similar lengths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_val_test_dataset = ClickBaitDataset()\n",
    "\n",
    "eval_size = math.floor(0.1 * len(train_val_test_dataset))\n",
    "train_val_test_dataset_indexes = list(range(len(train_val_test_dataset)))\n",
    "train_split_indexes = train_val_test_dataset_indexes[2 * eval_size:]\n",
    "val_split_indexes = train_val_test_dataset_indexes[eval_size:2 * eval_size]\n",
    "test_split_indexes = train_val_test_dataset_indexes[:eval_size]\n",
    "\n",
    "batch_size = 32\n",
    "train_dataloader = DataLoader(\n",
    "    train_val_test_dataset,\n",
    "    sampler=SubsetOrderedBatchWiseRandomSampler(\n",
    "        train_split_indexes,\n",
    "        get_order_value_callable=lambda example_index: train_val_test_dataset[example_index][1][1],\n",
    "        batch_size=batch_size\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=ClickBaitDataset.collate_fn\n",
    ")\n",
    "\n",
    "val_dataloader = DataLoader(\n",
    "    train_val_test_dataset,\n",
    "    sampler=SubsetOrderedSequentialSampler(\n",
    "        val_split_indexes,\n",
    "        get_order_value_callable=lambda example_index: train_val_test_dataset[example_index][1][1]\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=ClickBaitDataset.collate_fn\n",
    ")\n",
    "\n",
    "test_dataloader = DataLoader(\n",
    "    train_val_test_dataset,\n",
    "    sampler=SubsetOrderedSequentialSampler(\n",
    "        test_split_indexes,\n",
    "        get_order_value_callable=lambda example_index: train_val_test_dataset[example_index][1][1]\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=ClickBaitDataset.collate_fn\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we create the model and we wrap it with a `System` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BERTModel()\n",
    "\n",
    "last_activation = nn.Sigmoid()\n",
    "if torch.cuda.is_available():\n",
    "    system = System(model, last_activation=last_activation, device=torch.device('cuda'))\n",
    "else:\n",
    "    system = System(model, last_activation=last_activation, device=torch.device('cpu'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we finetune the model on the training set, using a small learning rate (0.00005)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_wrapper = GenericPointWiseLossWrapper(nn.BCEWithLogitsLoss())\n",
    "evals = {\n",
    "\n",
    "    'acc': evaluators.AccuracyEvaluator(),\n",
    "    'f1': evaluators.F1Evaluator(),\n",
    "    'auc': evaluators.AUROCEvaluator()\n",
    "\n",
    "}\n",
    "\n",
    "optimizer = torch.optim.Adam(\n",
    "    filter(lambda p: p.requires_grad, system.model.parameters()),\n",
    "    lr=0.00005\n",
    ")\n",
    "\n",
    "_ = system.train(\n",
    "    loss_wrapper,\n",
    "    optimizer,\n",
    "    train_data_loader=train_dataloader,\n",
    "    evaluators=evals,\n",
    "    evaluation_data_loaders={\n",
    "        'val': val_dataloader\n",
    "    },\n",
    "    callbacks=[\n",
    "        EarlyStoppingCriterionCallback(\n",
    "            patience=3,\n",
    "            evaluation_data_loader_key='val',\n",
    "            evaluator_key='f1',\n",
    "            tmp_best_state_filepath='data/click_bait_cur_best.weights'\n",
    "        )\n",
    "    ]\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we evaluate the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = system.evaluate(test_dataloader, evals)\n",
    "for r in results:\n",
    "    print(results[r])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the `predict` method in order to predict for all the examples returned by a `Dataloder`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = system.predict(test_dataloader, perform_last_activation=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_id = 3\n",
    "input_loc = 1\n",
    "text_loc = 0\n",
    "\n",
    "print(bert_tokenizer.decode(train_val_test_dataset[example_id][input_loc][text_loc]))\n",
    "print(predictions['outputs'][example_id])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we save the model's weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system.save_model_state('data/click_bait_final.weights')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}