01_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import logging\n",
    "logging.disable(logging.CRITICAL)\n",
    "import torch\n",
    "from transformers import DistilBertModel, DistilBertForSequenceClassification, DistilBertConfig\n",
    "from emotion_transformer.dataloader import dataloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model\n",
    "\n",
    "> construction of the DoubleDistilBert model for the SemEval-2019 Task 3 dataset (contextual emotion detection in text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer Sentence Embeddings\n",
    "\n",
    "First we create sentence embeddings for each utterance. We use a pretrained DistilBert model to obtain contextual word embeddings and then concatenate the CLS token embedding and the mean of the last layer. Note that in order to feed batches into out model we need to temporarily flatten our `input_ids`, i.e. we get three times as many input sentences as the specified `batch_size`.\n",
    "\n",
    "For more information on the (Distil)Bert models one can look at \n",
    "Jay Alammar's blog posts ([A Visual Guide to Using BERT for the First Time](https://jalammar.github.io/a-visual-guide-to-using-bert-for-the-first-time/) and [The Illustrated BERT, ELMo, and co.](https://jalammar.github.io/illustrated-bert/)) where also the following illustration is taken from.\n",
    "\n",
    "![DistilBert output](./images/bert-distilbert-output-tensor-predictions.png)\n",
    "\n",
    "Further references:\n",
    " \n",
    "* [DistilBert paper](https://arxiv.org/abs/1910.01108) and [blog post](https://medium.com/huggingface/distilbert-8cf3380435b5)\n",
    "* [Original Bert (Bidirectional Encoder Representations from Transformers) paper](https://arxiv.org/abs/1810.04805)\n",
    "* [tutorial for custom PyTorch modules](https://pytorch.org/tutorials/beginner/pytorch_with_examples.html)\n",
    "* [Huggingface transformers documentation](https://huggingface.co/transformers/v2.3.0/index.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class sentence_embeds_model(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    instantiates the pretrained DistilBert model and the linear layer\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dropout = 0.1):\n",
    "        super(sentence_embeds_model, self).__init__()\n",
    "        \n",
    "        self.transformer = DistilBertModel.from_pretrained('distilbert-base-uncased', dropout=dropout, \n",
    "                                                           output_hidden_states=True)\n",
    "        self.embedding_size = 2 * self.transformer.config.hidden_size\n",
    "        \n",
    "    def layerwise_lr(self, lr, decay):\n",
    "        \"\"\"\n",
    "        returns grouped model parameters with layer-wise decaying learning rate\n",
    "        \"\"\"\n",
    "        bert = self.transformer\n",
    "        num_layers = bert.config.n_layers\n",
    "        opt_parameters = [{'params': bert.embeddings.parameters(), 'lr': lr*decay**num_layers}]\n",
    "        opt_parameters += [{'params': bert.transformer.layer[l].parameters(), 'lr': lr*decay**(num_layers-l+1)} \n",
    "                            for l in range(num_layers)]\n",
    "        return opt_parameters\n",
    "               \n",
    "    def forward(self, input_ids = None, attention_mask = None, input_embeds = None):\n",
    "        \"\"\"\n",
    "        returns the sentence embeddings\n",
    "        \"\"\"\n",
    "        if input_ids is not None:\n",
    "            input_ids = input_ids.flatten(end_dim = 1)\n",
    "        if attention_mask is not None:\n",
    "            attention_mask = attention_mask.flatten(end_dim = 1)\n",
    "        output = self.transformer(input_ids = input_ids, \n",
    "                                  attention_mask = attention_mask, inputs_embeds = input_embeds)\n",
    "    \n",
    "        cls = output[0][:,0]\n",
    "        hidden_mean = torch.mean(output[1][-1],1)\n",
    "        sentence_embeds = torch.cat([cls, hidden_mean], dim = -1)\n",
    "        \n",
    "        return sentence_embeds.view(-1, 3, self.embedding_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To illustrate the model let us import our dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = 'data/clean_train.txt'\n",
    "batch_size = 5\n",
    "max_seq_len = 10\n",
    "emo_dict = {'others': 0, 'sad': 1, 'angry': 2, 'happy': 3}\n",
    "loader = dataloader(path, max_seq_len, batch_size, emo_dict)\n",
    "input_ids, attention_mask, labels = next(iter(loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The DistilBert model outputs \n",
    "\n",
    "* 768-dimensional embeddings for each of the 'max_seq_len' tokens and each of the three utterances of the `batch_size` conversations and\n",
    "* a list of the hidden-states in all of the 6 DistilBert transformer layers (including the first embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7, torch.Size([15, 10, 768]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeds_model = sentence_embeds_model()\n",
    "\n",
    "last_layer, hidden_states = embeds_model.transformer(input_ids = input_ids.flatten(end_dim = 1), attention_mask = attention_mask.flatten(end_dim = 1))\n",
    "input_embeds = embeds_model.transformer.embeddings(input_ids.flatten(end_dim = 1))\n",
    "\n",
    "assert torch.all(hidden_states[0] == input_embeds)\n",
    "assert torch.all(hidden_states[-1] == last_layer)\n",
    "\n",
    "len(hidden_states), last_layer.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now create sentence embeddings (we put the model in evaluation mode to deactivate dropout for later consistency checks). Note that the forward method of the model reshapes the output again back to the shape of the corresponding `input_ids`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeds_model.eval()\n",
    "assert(embeds_model.transformer.transformer.layer[0].dropout.training == False)\n",
    "\n",
    "sentence_embeds = embeds_model(input_ids = input_ids, attention_mask = attention_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 3, 10]), torch.Size([5, 3, 1536]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert input_ids.shape[:2] == sentence_embeds.shape[:2]\n",
    "assert sentence_embeds.shape[-1] == embeds_model.embedding_size\n",
    "input_ids.shape, sentence_embeds.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also check if the `layerwise_lr` method outputs all model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "for group in embeds_model.layerwise_lr(2.0e-5,0.95):\n",
    "    count += len(list(group['params']))\n",
    "\n",
    "assert count == len(list(embeds_model.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Context Transformer and Classification\n",
    "\n",
    "Next we use another transformer model to create contextual sentence embeddings, i.e. we model that a conversation consists of three utterances. This is partly motivated by the [BERTSUM paper](https://arxiv.org/abs/1903.10318).\n",
    "\n",
    "Moreover, we add a classification model for the emotion of the last utterance where we augment the loss by a binary loss due to the unbalanced data.\n",
    "\n",
    "Note that for our convenience we use\n",
    "\n",
    "* a linear projection of the sentence embeddings to a given `projection_size`\n",
    "* a (not pre-trained) DistilBertForSequenceClassification and flip the order of the utterances as the first input embedding gets classified by default\n",
    "* only one attention head, see also the paper [Are Sixteen Heads Really Better than One?](https://arxiv.org/abs/1905.10650)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "class context_classifier_model(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    instantiates the DisitlBertForSequenceClassification model, the position embeddings of the utterances, \n",
    "    and the binary loss function\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, embedding_size, projection_size, n_layers, emo_dict, dropout = 0.1):\n",
    "        super(context_classifier_model, self).__init__()\n",
    "        \n",
    "        self.projection_size = projection_size\n",
    "        self.projection = torch.nn.Linear(embedding_size, projection_size)         \n",
    "        self.position_embeds = torch.nn.Embedding(3, projection_size)\n",
    "        self.norm = torch.nn.LayerNorm(projection_size)\n",
    "        self.drop = torch.nn.Dropout(dropout)\n",
    "    \n",
    "        context_config = DistilBertConfig(dropout=dropout, \n",
    "                                dim=projection_size,\n",
    "                                hidden_dim=4*projection_size,\n",
    "                                n_layers=n_layers,\n",
    "                                n_heads = 1,\n",
    "                                num_labels=4)\n",
    "\n",
    "        self.context_transformer = DistilBertForSequenceClassification(context_config)\n",
    "        self.others_label = emo_dict['others']\n",
    "        self.bin_loss_fct = torch.nn.BCEWithLogitsLoss()\n",
    "        \n",
    "    def bin_loss(self, logits, labels):\n",
    "        \"\"\"\n",
    "        defined the additional binary loss for the `others` label\n",
    "        \"\"\"\n",
    "        bin_labels = torch.where(labels == self.others_label, torch.ones_like(labels), \n",
    "                                 torch.zeros_like(labels)).float()\n",
    "        bin_logits = logits[:, self.others_label]    \n",
    "        return self.bin_loss_fct(bin_logits, bin_labels)\n",
    "\n",
    "    def forward(self, sentence_embeds, labels = None):\n",
    "        \"\"\"\n",
    "        returns the logits and the corresponding loss if `labels` are given\n",
    "        \"\"\"\n",
    "        \n",
    "        position_ids = torch.arange(3, dtype=torch.long, device=sentence_embeds.device)\n",
    "        position_ids = position_ids.expand(sentence_embeds.shape[:2]) \n",
    "        position_embeds = self.position_embeds(position_ids)\n",
    "        sentence_embeds = self.projection(sentence_embeds) + position_embeds \n",
    "        sentence_embeds = self.drop(self.norm(sentence_embeds))\n",
    "        if labels is None:\n",
    "            return self.context_transformer(inputs_embeds = sentence_embeds.flip(1), labels = labels)[0]\n",
    "        \n",
    "        else:\n",
    "            loss, logits = self.context_transformer(inputs_embeds = sentence_embeds.flip(1), labels = labels)\n",
    "            return loss + self.bin_loss(logits, labels), logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us initiate a the `context_classifier_model` with the corresponding `projection_size` of the sentence embedding model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "projection_size = 100\n",
    "n_layers = 2\n",
    "\n",
    "classifier = context_classifier_model(embeds_model.embedding_size, projection_size, n_layers, emo_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and do some basic checks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(2.1022, grad_fn=<AddBackward0>),\n",
       " tensor([[-0.0509, -0.0632, -0.0165,  0.0358],\n",
       "         [-0.0536, -0.0644, -0.0126,  0.0392],\n",
       "         [-0.0487, -0.0654, -0.0186,  0.0388],\n",
       "         [-0.0541, -0.0641, -0.0206,  0.0491],\n",
       "         [-0.0566, -0.0647, -0.0171,  0.0375]], grad_fn=<AddmmBackward>))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier.eval()\n",
    "assert(classifier.context_transformer.distilbert.transformer.layer[0].dropout.training == False)\n",
    "\n",
    "loss, logits = classifier(sentence_embeds, labels = labels)\n",
    "assert torch.all(logits == classifier(sentence_embeds))\n",
    "assert loss == torch.nn.CrossEntropyLoss()(logits, labels) + classifier.bin_loss(logits, labels)\n",
    "\n",
    "loss, logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, for our main consistency check we compute the gradient of the loss w.r.t. to the input embeddings. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_embeds = input_embeds.clone().detach().requires_grad_(True)\n",
    "sentence_embeds_check = embeds_model(input_embeds = input_embeds, attention_mask = attention_mask)\n",
    "logits_check = classifier(sentence_embeds_check)\n",
    "assert torch.all(logits == logits_check)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0000e+00,  0.0000e+00,  0.0000e+00, -7.6451e-07, -1.2572e-06,\n",
       "        -7.2055e-05,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
       "         0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits_check[1,0].backward()\n",
    "input_embeds.grad[:,0,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As anticipated, we see that only the fourth, fifth, and sixth input embedding effect the second prediction.\n",
    "These correspond to the second conversation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.all(input_embeds[3:6] == embeds_model.transformer.embeddings(input_ids[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we define the metrics, i.e. microaveraged precision, recall, and f1-score (ignoring the others class), for the evaluation of our model according to the [SemEval-2019 Task 3 challenge](https://www.aclweb.org/anthology/S19-2005/). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def metrics(loss, logits, labels):\n",
    "    cm = torch.zeros((4,4), device = loss.device)\n",
    "    preds = torch.argmax(logits, dim=1)\n",
    "    acc = (labels == preds).float().mean()\n",
    "    for label, pred in zip(labels.view(-1), preds.view(-1)):\n",
    "        cm[label.long(), pred.long()] += 1\n",
    "        \n",
    "    tp = cm.diagonal()[1:].sum()\n",
    "    fp = cm[:, 1:].sum() - tp\n",
    "    fn = cm[1:, :].sum() - tp \n",
    "    return {'val_loss': loss, 'val_acc': acc, 'tp': tp, 'fp': fp, 'fn': fn}\n",
    "\n",
    "def f1_score(tp, fp, fn):\n",
    "    prec_rec_f1 = {}\n",
    "    prec_rec_f1['precision'] = tp / (tp + fp)\n",
    "    prec_rec_f1['recall'] = tp / (tp + fn)\n",
    "    prec_rec_f1['f1_score'] = 2 * (prec_rec_f1['precision'] * prec_rec_f1['recall']) / (prec_rec_f1['precision'] + prec_rec_f1['recall'])\n",
    "    return prec_rec_f1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'val_loss': tensor(2.1022, grad_fn=<AddBackward0>),\n",
       " 'val_acc': tensor(0.),\n",
       " 'tp': tensor(0.),\n",
       " 'fp': tensor(5.),\n",
       " 'fn': tensor(3.)}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metric = metrics(loss, logits, labels)\n",
    "metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'precision': tensor(0.), 'recall': tensor(0.), 'f1_score': tensor(nan)}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_score(metric['tp'], metric['fp'], metric['fn'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}