lessons/5-NLP/18-Transformers/TransformersPyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention mechanisms and transformers\n",
    "\n",
    "One major drawback of recurrent networks is that all words in a sequence have the same impact on the result. This causes sub-optimal performance with standard LSTM encoder-decoder models for sequence to sequence tasks, such as Named Entity Recognition and Machine Translation. In reality specific words in the input sequence often have more impact on sequential outputs than others.\n",
    "\n",
    "Consider sequence-to-sequence model, such as machine translation. It is implemented by two recurrent networks, where one network (**encoder**) would collapse input sequence into hidden state, and another one, **decoder**, would unroll this hidden state into translated result. The problem with this approach is that final state of the network would have hard time remembering the beginning of a sentence, thus causing poor quality of the model on long sentences.\n",
    "\n",
    "**Attention Mechanisms** provide a means of weighting the contextual impact of each input vector on each output prediction of the RNN. The way it is implemented is by creating shortcuts between intermediate states of the input RNN, and output RNN. In this manner, when generating output symbol $y_t$, we will take into account all input hidden states $h_i$, with different weight coefficients $\\alpha_{t,i}$. \n",
    "\n",
    "![Image showing an encoder/decoder model with an additive attention layer](./images/encoder-decoder-attention.png)\n",
    "*The encoder-decoder model with additive attention mechanism in [Bahdanau et al., 2015](https://arxiv.org/pdf/1409.0473.pdf), cited from [this blog post](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html)*\n",
    "\n",
    "Attention matrix $\\{\\alpha_{i,j}\\}$ would represent the degree which certain input words play in generation of a given word in the output sequence. Below is the example of such a matrix:\n",
    "\n",
    "![Image showing a sample alignment found by RNNsearch-50, taken from Bahdanau - arviz.org](./images/bahdanau-fig3.png)\n",
    "\n",
    "*Figure taken from [Bahdanau et al., 2015](https://arxiv.org/pdf/1409.0473.pdf) (Fig.3)*\n",
    "\n",
    "Attention mechanisms are responsible for much of the current or near current state of the art in Natural language processing. Adding attention however greatly increases the number of model parameters which led to scaling issues with RNNs. A key constraint of scaling RNNs is that the recurrent nature of the models makes it challenging to batch and parallelize training. In an RNN each element of a sequence needs to be processed in sequential order which means it cannot be easily parallelized.\n",
    "\n",
    "Adoption of attention mechanisms combined with this constraint led to the creation of the now State of the Art Transformer Models that we know and use today from BERT to OpenGPT3.\n",
    "\n",
    "## Transformer models\n",
    "\n",
    "Instead of forwarding the context of each previous prediction into the next evaluation step, **transformer models** use **positional encodings** and attention to capture the context of a given input with in a provided window of text. The image below shows how positional encodings with attention can capture context within a given window.\n",
    "\n",
    "![Animated GIF showing how the evaluations are performed in transformer models.](./images/transformer-animated-explanation.gif) \n",
    "\n",
    "Since each input position is mapped independently to each output position, transformers can parallelize better than RNNs, which enables much larger and more expressive language models. Each attention head can be used to learn different relationships between words that improves downstream Natural Language Processing tasks.\n",
    "\n",
    "**BERT** (Bidirectional Encoder Representations from Transformers) is a very large multi layer transformer network with 12 layers for *BERT-base*, and 24 for *BERT-large*. The model is first pre-trained on large corpus of text data (WikiPedia + books) using unsupervised training (predicting masked words in a sentence). During pre-training the model absorbs significant level of language understanding which can then be leveraged with other datasets using fine tuning. This process is called **transfer learning**. \n",
    "\n",
    "![picture from http://jalammar.github.io/illustrated-bert/](./images/jalammarBERT-language-modeling-masked-lm.png)\n",
    "\n",
    "There are many variations of Transformer architectures including BERT, DistilBERT. BigBird, OpenGPT3 and more that can be fine tuned. The [HuggingFace package](https://github.com/huggingface/) provides repository for training many of these architectures with PyTorch. \n",
    "\n",
    "## Using BERT for text classification\n",
    "\n",
    "Let's see how we can use pre-trained BERT model for solving our traditional task: sequence classification. We will classify our original AG News dataset.\n",
    "\n",
    "First, let's load HuggingFace library and our dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading dataset...\n",
      "Building vocab...\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchtext\n",
    "from torchnlp import *\n",
    "import transformers\n",
    "train_dataset, test_dataset, classes, vocab = load_dataset()\n",
    "vocab_len = len(vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we will be using pre-trained BERT model, we would need to use specific tokenizer. First, we will load a tokenizer associated with pre-trained BERT model.\n",
    "\n",
    "HuggingFace library contains a repository of pre-trained models, which you can use just by specifying their names as arguments to `from_pretrained` functions. All required binary files for the model would automatically be downloaded.\n",
    "\n",
    "However, at certain times you would need to load your own models, in which case you can specify the directory that contains all relevant files, including parameters for tokenizer, `config.json` file with model parameters, binary weights, etc. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To load the model from Internet repository using model name. \n",
    "# Use this if you are running from your own copy of the notebooks\n",
    "bert_model = 'bert-base-uncased' \n",
    "\n",
    "# To load the model from the directory on disk. Use this for Microsoft Learn module, because we have\n",
    "# prepared all required files for you.\n",
    "bert_model = './bert'\n",
    "\n",
    "tokenizer = transformers.BertTokenizer.from_pretrained(bert_model)\n",
    "\n",
    "MAX_SEQ_LEN = 128\n",
    "PAD_INDEX = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)\n",
    "UNK_INDEX = tokenizer.convert_tokens_to_ids(tokenizer.unk_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `tokenizer` object contains the `encode` function that can be directly used to encode text:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[101, 1052, 22123, 2953, 2818, 2003, 1037, 2307, 7705, 2005, 17953, 2361, 102]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.encode('PyTorch is a great framework for NLP')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, let's create iterators which we will use during training to access the data. Because BERT uses it's own encoding function, we would need to define a padding function similar to `padify` we have defined before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_bert(b):\n",
    "    # b is the list of tuples of length batch_size\n",
    "    #   - first element of a tuple = label, \n",
    "    #   - second = feature (text sequence)\n",
    "    # build vectorized sequence\n",
    "    v = [tokenizer.encode(x[1]) for x in b]\n",
    "    # compute max length of a sequence in this minibatch\n",
    "    l = max(map(len,v))\n",
    "    return ( # tuple of two tensors - labels and features\n",
    "        torch.LongTensor([t[0] for t in b]),\n",
    "        torch.stack([torch.nn.functional.pad(torch.tensor(t),(0,l-len(t)),mode='constant',value=0) for t in v])\n",
    "    )\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, collate_fn=pad_bert, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, collate_fn=pad_bert)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our case, we will be using pre-trained BERT model called `bert-base-uncased`. Let's load the model using `BertForSequenceClassfication` package. This ensures that our model already has a required architecture for classification, including final classifier. You will see warning message stating that weights of the final classifier are not initialized, and model would require pre-training - that is perfectly okay, because it is exactly what we are about to do!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at ./bert were not used when initializing BertForSequenceClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at ./bert and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = transformers.BertForSequenceClassification.from_pretrained(bert_model,num_labels=4).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to begin training! Because BERT is already pre-trained, we want to start with rather small learning rate in order not to destroy initial weights.\n",
    "\n",
    "All hard work is done by `BertForSequenceClassification` model. When we call the model on the training data, it returns both loss and network output for input minibatch. We use loss for parameter optimization (`loss.backward()` does the backward pass), and `out` for computing training accuracy by comparing obtained labels `labs` (computed using `argmax`) with expected `labels`.\n",
    "\n",
    "In order to control the process, we accumulate loss and accuracy over several iterations, and print them every `report_freq` training cycles.\n",
    "\n",
    "This training will likely take quite a long time, so we limit the number of iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss = 1.1254194641113282, Accuracy = 0.585\n",
      "Loss = 0.6194715118408203, Accuracy = 0.83\n",
      "Loss = 0.46665248870849607, Accuracy = 0.8475\n",
      "Loss = 0.4309701919555664, Accuracy = 0.8575\n",
      "Loss = 0.35427074432373046, Accuracy = 0.8825\n",
      "Loss = 0.3306886291503906, Accuracy = 0.8975\n",
      "Loss = 0.30340143203735354, Accuracy = 0.8975\n",
      "Loss = 0.26139299392700194, Accuracy = 0.915\n",
      "Loss = 0.26708646774291994, Accuracy = 0.9225\n",
      "Loss = 0.3667240524291992, Accuracy = 0.8675\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=2e-5)\n",
    "\n",
    "report_freq = 50\n",
    "iterations = 500 # make this larger to train for longer time!\n",
    "\n",
    "model.train()\n",
    "\n",
    "i,c = 0,0\n",
    "acc_loss = 0\n",
    "acc_acc = 0\n",
    "\n",
    "for labels,texts in train_loader:\n",
    "    labels = labels.to(device)-1 # get labels in the range 0-3         \n",
    "    texts = texts.to(device)\n",
    "    loss, out = model(texts, labels=labels)[:2]\n",
    "    labs = out.argmax(dim=1)\n",
    "    acc = torch.mean((labs==labels).type(torch.float32))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    acc_loss += loss\n",
    "    acc_acc += acc\n",
    "    i+=1\n",
    "    c+=1\n",
    "    if i%report_freq==0:\n",
    "        print(f\"Loss = {acc_loss.item()/c}, Accuracy = {acc_acc.item()/c}\")\n",
    "        c = 0\n",
    "        acc_loss = 0\n",
    "        acc_acc = 0\n",
    "    iterations-=1\n",
    "    if not iterations:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see (especially if you increase the number of iterations and wait long enough) that BERT classification gives us pretty good accuracy! That is because BERT already understands quite well the structure of the language, and we only need to fine-tune final classifier. However, because BERT is a large model, the whole training process takes a long time, and requires serious computational power! (GPU, and preferably more than one).\n",
    "\n",
    "> **Note:** In our example, we have been using one of the smallest pre-trained BERT models. There are larger models that are likely to yield better results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating the model performance\n",
    "\n",
    "Now we can evaluate performance of our model on test dataset. Evaluation loop is pretty similar to training loop, but we should not forget to switch model to evaluation mode by calling `model.eval()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final accuracy: 0.9047029702970297\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "iterations = 100\n",
    "acc = 0\n",
    "i = 0\n",
    "for labels,texts in test_loader:\n",
    "    labels = labels.to(device)-1      \n",
    "    texts = texts.to(device)\n",
    "    _, out = model(texts, labels=labels)[:2]\n",
    "    labs = out.argmax(dim=1)\n",
    "    acc += torch.mean((labs==labels).type(torch.float32))\n",
    "    i+=1\n",
    "    if i>iterations: break\n",
    "        \n",
    "print(f\"Final accuracy: {acc.item()/i}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Takeaway\n",
    "\n",
    "In this unit, we have seen how easy it is to take pre-trained language model from **transformers** library and adapt it to our text classification task. Similarly, BERT models can be used for entity extraction, question answering, and other NLP tasks.\n",
    "\n",
    "Transformer models represent current state-of-the-art in NLP, and in most of the cases it should be the first solution you start experimenting with when implementing custom NLP solutions. However, understanding basic underlying principles of recurrent neural networks discussed in this module is extremely important if you want to build advanced neural models."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py37_pytorch",
   "language": "python",
   "name": "conda-env-py37_pytorch-py"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}