week-7/7-Deep-Learning-and-Transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Week 7- Accounting for Context (Deep Learning & Transformers)\n",
    "\n",
    "In this week we introduce the use of Deep Neural Networks to work with text. These models make use of word embeddings to learn complex semantic and syntactic relations. We focus on BERT, which is a language model built using bidirectional encoders. It allows us to take advantage of a powerful pre-trained model which we can then use to perform our own tasks based on data we analyze.\n",
    "\n",
    "In this notebook we use ```huggingface/transformers```, a Python package that allows for easy interface to use pre-trained BERT models. It is built using Tensorflow and PyTorch, two computational graph packages which are built specifically for creating powerful neural networks. We will also be introducing Keras, which allows us to easily build Neural Networks in an abstracted way. Keras is a popular way to understand how we can stack layers to create such Neural Networks, but to reach state-of-the-art results we will stick with using BERT and similar models that can be tuned to extremely high performance on specific language understanding and generation tasks.\n",
    "\n",
    "To demonstrate this, we begin by using the [Corpus of Linguistic Acceptability](https://nyu-mll.github.io/CoLA/). We will also use BERT by learning how to extract embeddings from such a model and use it to semantically probe sentences.\n",
    "\n",
    "## NOTE\n",
    "\n",
    "This notebook **requires** GPUs for training models in section 1 and section 3. To train models, please use [the Google Colab version of this notebook](https://colab.research.google.com/drive/1LMFV9oUGLo68turYY5QBXWKOadRqAAZ0) to create the models, unless you have a suitable GPU and software to do it on your local machine.\n",
    "\n",
    "If you run the computationally intensive models on a CPU, they will take a long time!\n",
    "\n",
    "Also, while this is one of the longer assignments, much of the Jupyter Notebook is neural network code that you can read through quickly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch # pip install torch\n",
    "from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n",
    "from sklearn.model_selection import train_test_split\n",
    "from transformers import BertTokenizer, BertConfig # pip install tranformers\n",
    "from transformers import AdamW, BertForSequenceClassification\n",
    "from tqdm import tqdm, trange\n",
    "import pandas as pd\n",
    "import io\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import lucem_illud #pip install -U git+https://github.com/UChicago-Computational-Content-Analysis/lucem_illud.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KERAS_BACKEND'] = 'tensorflow'\n",
    "from keras.preprocessing.sequence import pad_sequences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CoLA Dataset and pre-processing\n",
    "\n",
    "We start with loading our dataset and pre-processing it. The pre-processing follows similar steps as we have done in the past, but we will be using pre-written modules offered by the transformers package. These are some of the things we have to take care of when using this particular BERT model.\n",
    "\n",
    "    -special tokens to mark the beginning ([CLS]) and separation/end of sentences ([SEP])\n",
    "    -tokens that conforms with the fixed vocabulary used in BERT\n",
    "    -token IDs from BERT’s tokenizer\n",
    "    -mask IDs to indicate which elements in the sequence are tokens and which are padding elements\n",
    "    -segment IDs used to distinguish different sentences\n",
    "    -positional embeddings used to show token position within the sequence\n",
    "\n",
    "\n",
    "We will be using parts of the code from [this notebook](https://colab.research.google.com/drive/1ywsvwO6thOVOrfagjjfuxEf6xVRxbUNO#scrollTo=BJR6t_gCQe_x) which walks us through the process of using a pre-trained BERT model. The interface to use these models comes from the package [huggingface/transformers](https://github.com/huggingface/transformers). \n",
    "\n",
    "We start by setting up our GPU if we can - this may not work on your machine, so it has been commented out.\n",
    "\n",
    "An aside: check out this tutorial too - https://mccormickml.com/2019/07/22/BERT-fine-tuning/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "if gpu:\n",
    "    n_gpu = torch.cuda.device_count()\n",
    "    torch.cuda.get_device_name(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"../data/cola_public/raw/in_domain_train.tsv\", delimiter='\\t', header=None, names=['sentence_source', 'label', 'label_notes', 'sentence'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create sentence and label lists\n",
    "sentences = df.sentence.values\n",
    "\n",
    "# We need to add special tokens at the beginning and end of each sentence for BERT to work properly\n",
    "sentences = [\"[CLS] \" + sentence + \" [SEP]\" for sentence in sentences]\n",
    "labels = df.label.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)\n",
    "\n",
    "tokenized_texts = [tokenizer.tokenize(sent) for sent in sentences]\n",
    "print (\"Tokenize the first sentence:\")\n",
    "print (tokenized_texts[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the maximum sequence length. The longest sequence in our training set is 47, but we'll leave room on the end anyway. \n",
    "# In the original paper, the authors used a length of 512.\n",
    "MAX_LEN = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the BERT tokenizer to convert the tokens to their index numbers in the BERT vocabulary\n",
    "input_ids = [tokenizer.convert_tokens_to_ids(x) for x in tokenized_texts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pad our input tokens\n",
    "input_ids = pad_sequences(input_ids, maxlen=MAX_LEN, dtype=\"long\", truncating=\"post\", padding=\"post\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create attention masks\n",
    "attention_masks = []\n",
    "\n",
    "# Create a mask of 1s for each token followed by 0s for padding\n",
    "for seq in input_ids:\n",
    "    seq_mask = [float(i>0) for i in seq]\n",
    "    attention_masks.append(seq_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use train_test_split to split our data into train and validation sets for training\n",
    "train_inputs, validation_inputs, train_labels, validation_labels = train_test_split(input_ids, labels, \n",
    "                                                            random_state=2020, test_size=0.1)\n",
    "train_masks, validation_masks, _, _ = train_test_split(attention_masks, input_ids,\n",
    "                                             random_state=2020, test_size=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Introducing Deep Neural Nets with Keras LSTM\n",
    "\n",
    "A popular, simplified package for introducing deep neural networks is [Keras](https://keras.io). It is a high level package in that we don't bother with every detail or hyper-parameter associated with the neural network (e.g., regularizers), and can stack on layers directly. For a rapid tutorial on neural networks for text such as the LSTM or the Recurrent Neural Network, Colah's blog is a great start. [LSTM](https://colah.github.io/posts/2015-08-Understanding-LSTMs/) is an article on LSTMs, and if you'd like to  learn about RNN, Andrej Karpathy does a great job in [this blog post](http://karpathy.github.io/2015/05/21/rnn-effectiveness/), in addition to our reading from the newest online verion of Jurafsky & Martin's review of deep learning methods in their book on speech and language processing, chapters [6,7,9,10](https://web.stanford.edu/~jurafsky/slp3/), and the [*Deep Learning*](https://www.deeplearningbook.org/) book by Goodfellow, Bengio & Courville.\n",
    "\n",
    "Before we get to BERT, in the following cells we build a basic deep net that has an embedding layer and an LSTM to perform classification. This is to illustrate the process of using Keras, which is a very popular library for such work. It may not yield state of the art performance because it constrains the hyperparameters you can tune, but is nonetheless an useful tool and works well on some datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.models import Sequential\n",
    "from keras.layers import Dense, LSTM, Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_in_size = tokenizer.vocab_size\n",
    "embedding_dim = 32\n",
    "unit = 100\n",
    "no_labels = len(np.unique(train_labels))\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_lstm = Sequential()\n",
    "model_lstm.add(Embedding(vocab_in_size, embedding_dim, input_length=MAX_LEN))\n",
    "model_lstm.add(LSTM(unit))\n",
    "model_lstm.add(Dense(no_labels, activation='softmax'))\n",
    "model_lstm.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])\n",
    "model_lstm.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history_lstm = model_lstm.fit(train_inputs, train_labels, \n",
    "                              epochs=10,batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the accuracy of this model isn't terrible, but it still hovers around 70%. Below there is code for a slightly modified neural network - how does this one perform? Note that in this model, I have added another LSTM layer. You are encouraged to explore the [Keras documentaion](https://keras.io/layers/about-keras-layers/) to explore what kind of layers you can add and how they change performances for different tasks. Different kinds of losses, optimizers, activations and layers can change the flavour of your net dramatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_lstm2 = Sequential()\n",
    "# model_lstm2.add(Embedding(vocab_in_size, embedding_dim, input_length=MAX_LEN))\n",
    "# model_lstm2.add(LSTM(units))\n",
    "# model_lstm2.add(LSTM(units))\n",
    "# model_lstm2.add(Dense(1, activation='sigmoid'))\n",
    "# model_lstm2.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])\n",
    "\n",
    "# history_lstm2 = model_lstm2.fit(input_data_train, labels, epochs=10, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CDTkQEaHNDe9"
   },
   "source": [
    "# Encoders, Decoders, Seq2Seq and Attention\n",
    "An important part of transformers are the Encoder and Decoder models, which are a powerful way to deal with sequence to sequence problems - this kind of setup is also called seq2seq.\n",
    "\n",
    "When we use this kind of sequence to sequence alignment and also map certain parts of the sequences with different values (a process called Attention), the sequence to sequence prediction task performance increases. In this section and the next we will focus on Machine Translation as the domain to try Encoder and Decoder tasks, without and then with attention. \n",
    "\n",
    "**We highly recommend using this [visual blog](http://jalammar.github.io/visualizing-neural-machine-translation-mechanics-of-seq2seq-models-with-attention/) which explains seq2seq and attention based models**\n",
    "\n",
    "We will be using the code from a Keras tutorial and PyTorch tutorial to demonstrate these models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xIuje1-Ve-67"
   },
   "source": [
    "## Character-level recurrent sequence-to-sequence model\n",
    "\n",
    "This code is adapted from the [Keras official tutorial](https://keras.io/examples/nlp/lstm_seq2seq/) on seq2seq using LSTM, authored by [fchollet](https://twitter.com/fchollet). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dW7DdSj0e-7J"
   },
   "source": [
    "### Introduction\n",
    "\n",
    "This example demonstrates how to implement a basic character-level\n",
    "recurrent sequence-to-sequence model. We apply it to translating\n",
    "short English sentences into short French sentences,\n",
    "character-by-character. Note that it is fairly unusual to\n",
    "do character-level machine translation, as word-level\n",
    "models are more common in this domain.\n",
    "\n",
    "**Summary of the algorithm**\n",
    "\n",
    "- We start with input sequences from a domain (e.g. English sentences)\n",
    "    and corresponding target sequences from another domain\n",
    "    (e.g. French sentences).\n",
    "- An encoder LSTM turns input sequences to 2 state vectors\n",
    "    (we keep the last LSTM state and discard the outputs).\n",
    "- A decoder LSTM is trained to turn the target sequences into\n",
    "    the same sequence but offset by one timestep in the future,\n",
    "    a training process called \"teacher forcing\" in this context.\n",
    "    It uses as initial state the state vectors from the encoder.\n",
    "    Effectively, the decoder learns to generate `targets[t+1...]`\n",
    "    given `targets[...t]`, conditioned on the input sequence.\n",
    "- In inference mode, when we want to decode unknown input sequences, we:\n",
    "    - Encode the input sequence into state vectors\n",
    "    - Start with a target sequence of size 1\n",
    "        (just the start-of-sequence character)\n",
    "    - Feed the state vectors and 1-char target sequence\n",
    "        to the decoder to produce predictions for the next character\n",
    "    - Sample the next character using these predictions\n",
    "        (we simply use argmax).\n",
    "    - Append the sampled character to the target sequence\n",
    "    - Repeat until we generate the end-of-sequence character or we\n",
    "        hit the character limit.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Uk8__4jEe-7L"
   },
   "source": [
    "### Setup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fjnD7Gtle-7M"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A9axnF-ie-7O"
   },
   "source": [
    "### Download the data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ik7TrAiMe-7P"
   },
   "outputs": [],
   "source": [
    "!!curl -O http://www.manythings.org/anki/fra-eng.zip\n",
    "!!unzip fra-eng.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OzpQ-1sre-7T"
   },
   "source": [
    "### Configuration\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "14m6qbZVe-7T"
   },
   "outputs": [],
   "source": [
    "batch_size = 64  # Batch size for training.\n",
    "epochs = 100  # Number of epochs to train for.\n",
    "latent_dim = 256  # Latent dimensionality of the encoding space.\n",
    "num_samples = 10000  # Number of samples to train on.\n",
    "# Path to the data txt file on disk.\n",
    "data_path = \"fra.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TJ1NzQs6e-7U"
   },
   "source": [
    "### Prepare the data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5VXgK_P4e-7V"
   },
   "outputs": [],
   "source": [
    "# Vectorize the data.\n",
    "input_texts = []\n",
    "target_texts = []\n",
    "input_characters = set()\n",
    "target_characters = set()\n",
    "with open(data_path, \"r\", encoding=\"utf-8\") as f:\n",
    "    lines = f.read().split(\"\\n\")\n",
    "for line in lines[: min(num_samples, len(lines) - 1)]:\n",
    "    input_text, target_text, _ = line.split(\"\\t\")\n",
    "    # We use \"tab\" as the \"start sequence\" character\n",
    "    # for the targets, and \"\\n\" as \"end sequence\" character.\n",
    "    target_text = \"\\t\" + target_text + \"\\n\"\n",
    "    input_texts.append(input_text)\n",
    "    target_texts.append(target_text)\n",
    "    for char in input_text:\n",
    "        if char not in input_characters:\n",
    "            input_characters.add(char)\n",
    "    for char in target_text:\n",
    "        if char not in target_characters:\n",
    "            target_characters.add(char)\n",
    "\n",
    "input_characters = sorted(list(input_characters))\n",
    "target_characters = sorted(list(target_characters))\n",
    "num_encoder_tokens = len(input_characters)\n",
    "num_decoder_tokens = len(target_characters)\n",
    "max_encoder_seq_length = max([len(txt) for txt in input_texts])\n",
    "max_decoder_seq_length = max([len(txt) for txt in target_texts])\n",
    "\n",
    "print(\"Number of samples:\", len(input_texts))\n",
    "print(\"Number of unique input tokens:\", num_encoder_tokens)\n",
    "print(\"Number of unique output tokens:\", num_decoder_tokens)\n",
    "print(\"Max sequence length for inputs:\", max_encoder_seq_length)\n",
    "print(\"Max sequence length for outputs:\", max_decoder_seq_length)\n",
    "\n",
    "input_token_index = dict([(char, i) for i, char in enumerate(input_characters)])\n",
    "target_token_index = dict([(char, i) for i, char in enumerate(target_characters)])\n",
    "\n",
    "encoder_input_data = np.zeros(\n",
    "    (len(input_texts), max_encoder_seq_length, num_encoder_tokens), dtype=\"float32\"\n",
    ")\n",
    "decoder_input_data = np.zeros(\n",
    "    (len(input_texts), max_decoder_seq_length, num_decoder_tokens), dtype=\"float32\"\n",
    ")\n",
    "decoder_target_data = np.zeros(\n",
    "    (len(input_texts), max_decoder_seq_length, num_decoder_tokens), dtype=\"float32\"\n",
    ")\n",
    "\n",
    "for i, (input_text, target_text) in enumerate(zip(input_texts, target_texts)):\n",
    "    for t, char in enumerate(input_text):\n",
    "        encoder_input_data[i, t, input_token_index[char]] = 1.0\n",
    "    encoder_input_data[i, t + 1 :, input_token_index[\" \"]] = 1.0\n",
    "    for t, char in enumerate(target_text):\n",
    "        # decoder_target_data is ahead of decoder_input_data by one timestep\n",
    "        decoder_input_data[i, t, target_token_index[char]] = 1.0\n",
    "        if t > 0:\n",
    "            # decoder_target_data will be ahead by one timestep\n",
    "            # and will not include the start character.\n",
    "            decoder_target_data[i, t - 1, target_token_index[char]] = 1.0\n",
    "    decoder_input_data[i, t + 1 :, target_token_index[\" \"]] = 1.0\n",
    "    decoder_target_data[i, t:, target_token_index[\" \"]] = 1.0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VmT-s5TGe-7X"
   },
   "source": [
    "### Build the model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b1emHhx2e-7Y"
   },
   "outputs": [],
   "source": [
    "# Define an input sequence and process it.\n",
    "encoder_inputs = keras.Input(shape=(None, num_encoder_tokens))\n",
    "encoder = keras.layers.LSTM(latent_dim, return_state=True)\n",
    "encoder_outputs, state_h, state_c = encoder(encoder_inputs)\n",
    "\n",
    "# We discard `encoder_outputs` and only keep the states.\n",
    "encoder_states = [state_h, state_c]\n",
    "\n",
    "# Set up the decoder, using `encoder_states` as initial state.\n",
    "decoder_inputs = keras.Input(shape=(None, num_decoder_tokens))\n",
    "\n",
    "# We set up our decoder to return full output sequences,\n",
    "# and to return internal states as well. We don't use the\n",
    "# return states in the training model, but we will use them in inference.\n",
    "decoder_lstm = keras.layers.LSTM(latent_dim, return_sequences=True, return_state=True)\n",
    "decoder_outputs, _, _ = decoder_lstm(decoder_inputs, initial_state=encoder_states)\n",
    "decoder_dense = keras.layers.Dense(num_decoder_tokens, activation=\"softmax\")\n",
    "decoder_outputs = decoder_dense(decoder_outputs)\n",
    "\n",
    "# Define the model that will turn\n",
    "# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`\n",
    "model = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6KVvUL8re-7b"
   },
   "source": [
    "### Train the model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wIXh70Mte-7c"
   },
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    optimizer=\"rmsprop\", loss=\"categorical_crossentropy\", metrics=[\"accuracy\"]\n",
    ")\n",
    "model.fit(\n",
    "    [encoder_input_data, decoder_input_data],\n",
    "    decoder_target_data,\n",
    "    batch_size=batch_size,\n",
    "    epochs=epochs,\n",
    "    validation_split=0.2,\n",
    ")\n",
    "# Save model\n",
    "model.save(\"s2s\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_HgC-DTCe-7d"
   },
   "source": [
    "### Run inference (sampling)\n",
    "\n",
    "1. encode input and retrieve initial decoder state\n",
    "2. run one step of decoder with this initial state\n",
    "and a \"start of sequence\" token as target.\n",
    "Output will be the next target token.\n",
    "3. Repeat with the current target token and current states\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2JSbU7pse-7f"
   },
   "outputs": [],
   "source": [
    "# Define sampling models\n",
    "# Restore the model and construct the encoder and decoder.\n",
    "model = keras.models.load_model(\"s2s\")\n",
    "\n",
    "encoder_inputs = model.input[0]  # input_1\n",
    "encoder_outputs, state_h_enc, state_c_enc = model.layers[2].output  # lstm_1\n",
    "encoder_states = [state_h_enc, state_c_enc]\n",
    "encoder_model = keras.Model(encoder_inputs, encoder_states)\n",
    "\n",
    "decoder_inputs = model.input[1]  # input_2\n",
    "decoder_state_input_h = keras.Input(shape=(latent_dim,), name=\"input_3\")\n",
    "decoder_state_input_c = keras.Input(shape=(latent_dim,), name=\"input_4\")\n",
    "decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]\n",
    "decoder_lstm = model.layers[3]\n",
    "decoder_outputs, state_h_dec, state_c_dec = decoder_lstm(\n",
    "    decoder_inputs, initial_state=decoder_states_inputs\n",
    ")\n",
    "decoder_states = [state_h_dec, state_c_dec]\n",
    "decoder_dense = model.layers[4]\n",
    "decoder_outputs = decoder_dense(decoder_outputs)\n",
    "decoder_model = keras.Model(\n",
    "    [decoder_inputs] + decoder_states_inputs, [decoder_outputs] + decoder_states\n",
    ")\n",
    "\n",
    "# Reverse-lookup token index to decode sequences back to\n",
    "# something readable.\n",
    "reverse_input_char_index = dict((i, char) for char, i in input_token_index.items())\n",
    "reverse_target_char_index = dict((i, char) for char, i in target_token_index.items())\n",
    "\n",
    "\n",
    "def decode_sequence(input_seq):\n",
    "    # Encode the input as state vectors.\n",
    "    states_value = encoder_model.predict(input_seq)\n",
    "\n",
    "    # Generate empty target sequence of length 1.\n",
    "    target_seq = np.zeros((1, 1, num_decoder_tokens))\n",
    "    # Populate the first character of target sequence with the start character.\n",
    "    target_seq[0, 0, target_token_index[\"\\t\"]] = 1.0\n",
    "\n",
    "    # Sampling loop for a batch of sequences\n",
    "    # (to simplify, here we assume a batch of size 1).\n",
    "    stop_condition = False\n",
    "    decoded_sentence = \"\"\n",
    "    while not stop_condition:\n",
    "        output_tokens, h, c = decoder_model.predict([target_seq] + states_value)\n",
    "\n",
    "        # Sample a token\n",
    "        sampled_token_index = np.argmax(output_tokens[0, -1, :])\n",
    "        sampled_char = reverse_target_char_index[sampled_token_index]\n",
    "        decoded_sentence += sampled_char\n",
    "\n",
    "        # Exit condition: either hit max length\n",
    "        # or find stop character.\n",
    "        if sampled_char == \"\\n\" or len(decoded_sentence) > max_decoder_seq_length:\n",
    "            stop_condition = True\n",
    "\n",
    "        # Update the target sequence (of length 1).\n",
    "        target_seq = np.zeros((1, 1, num_decoder_tokens))\n",
    "        target_seq[0, 0, sampled_token_index] = 1.0\n",
    "\n",
    "        # Update states\n",
    "        states_value = [h, c]\n",
    "    return decoded_sentence\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T8K7BU-fe-7i"
   },
   "source": [
    "You can now generate decoded sentences as such:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gHNiUUmPe-7j"
   },
   "outputs": [],
   "source": [
    "for seq_index in range(20):\n",
    "    # Take one sequence (part of the training set)\n",
    "    # for trying out decoding.\n",
    "    input_seq = encoder_input_data[seq_index : seq_index + 1]\n",
    "    decoded_sentence = decode_sequence(input_seq)\n",
    "    print(\"-\")\n",
    "    print(\"Input sentence:\", input_texts[seq_index])\n",
    "    print(\"Decoded sentence:\", decoded_sentence)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zTFsrjOgoTjB"
   },
   "source": [
    "\n",
    "##  Translation with a Sequence to Sequence Network and Attention\n",
    "\n",
    "This code is adapted from the [PyTorch official tutorial](https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html), authored by [Sean Robertson](https://github.com/spro/practical-pytorch).\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "03fSZ2pmv02A"
   },
   "source": [
    "### Introduction\n",
    "In this project we will be teaching a neural network to translate from\n",
    "French to English.\n",
    "\n",
    "\n",
    "    [KEY: > input, = target, < output]\n",
    "\n",
    "    > il est en train de peindre un tableau .\n",
    "    = he is painting a picture .\n",
    "    < he is painting a picture .\n",
    "\n",
    "    > pourquoi ne pas essayer ce vin delicieux ?\n",
    "    = why not try that delicious wine ?\n",
    "    < why not try that delicious wine ?\n",
    "\n",
    "    > elle n est pas poete mais romanciere .\n",
    "    = she is not a poet but a novelist .\n",
    "    < she not not a poet but a novelist .\n",
    "\n",
    "    > vous etes trop maigre .\n",
    "    = you re too skinny .\n",
    "    < you re all alone .\n",
    "\n",
    "... to varying degrees of success.\n",
    "\n",
    "This is made possible by the simple but powerful idea of the [sequence\n",
    "to sequence network](https://arxiv.org/abs/1409.3215), in which two\n",
    "recurrent neural networks work together to transform one sequence to\n",
    "another. An encoder network condenses an input sequence into a vector,\n",
    "and a decoder network unfolds that vector into a new sequence.\n",
    "\n",
    "\n",
    "To improve upon this model we'll use an [attention\n",
    "mechanism](https://arxiv.org/abs/1409.0473), which lets the decoder\n",
    "learn to focus over a specific range of the input sequence.\n",
    "\n",
    "### imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gIHXqUQvoTjD"
   },
   "outputs": [],
   "source": [
    "from __future__ import unicode_literals, print_function, division\n",
    "from io import open\n",
    "import unicodedata\n",
    "import string\n",
    "import re\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5xk9k0xgoTjE"
   },
   "source": [
    "### Loading data files\n",
    "\n",
    "\n",
    "The data for this project is a set of many thousands of English to\n",
    "French translation pairs.\n",
    "\n",
    "`This question on Open Data Stack\n",
    "Exchange <https://opendata.stackexchange.com/questions/3888/dataset-of-sentences-translated-into-many-languages>`__\n",
    "pointed me to the open translation site https://tatoeba.org/ which has\n",
    "downloads available at https://tatoeba.org/eng/downloads - and better\n",
    "yet, someone did the extra work of splitting language pairs into\n",
    "individual text files here: https://www.manythings.org/anki/\n",
    "\n",
    "The English to French pairs are too big to include in the repo, so\n",
    "download to ``data/eng-fra.txt`` before continuing. The file is a tab\n",
    "separated list of translation pairs:\n",
    "\n",
    "::\n",
    "\n",
    "    I am cold.    J'ai froid.\n",
    "\n",
    ".. Note::\n",
    "   Download the data from\n",
    "   `here <https://download.pytorch.org/tutorial/data.zip>`_\n",
    "   and extract it to the current directory.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F8VYnzaHoTjE"
   },
   "source": [
    "Similar to the character encoding used in the character-level RNN\n",
    "tutorials, we will be representing each word in a language as a one-hot\n",
    "vector, or giant vector of zeros except for a single one (at the index\n",
    "of the word). Compared to the dozens of characters that might exist in a\n",
    "language, there are many many more words, so the encoding vector is much\n",
    "larger. We will however cheat a bit and trim the data to only use a few\n",
    "thousand words per language.\n",
    "\n",
    ".. figure:: /_static/img/seq-seq-images/word-encoding.png\n",
    "   :alt:\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eUB-q_f8oTjG"
   },
   "source": [
    "We'll need a unique index per word to use as the inputs and targets of\n",
    "the networks later. To keep track of all this we will use a helper class\n",
    "called ``Lang`` which has word → index (``word2index``) and index → word\n",
    "(``index2word``) dictionaries, as well as a count of each word\n",
    "``word2count`` to use to later replace rare words.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yulUXQbioTjG"
   },
   "outputs": [],
   "source": [
    "SOS_token = 0\n",
    "EOS_token = 1\n",
    "\n",
    "\n",
    "class Lang:\n",
    "    def __init__(self, name):\n",
    "        self.name = name\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {0: \"SOS\", 1: \"EOS\"}\n",
    "        self.n_words = 2  # Count SOS and EOS\n",
    "\n",
    "    def addSentence(self, sentence):\n",
    "        for word in sentence.split(' '):\n",
    "            self.addWord(word)\n",
    "\n",
    "    def addWord(self, word):\n",
    "        if word not in self.word2index:\n",
    "            self.word2index[word] = self.n_words\n",
    "            self.word2count[word] = 1\n",
    "            self.index2word[self.n_words] = word\n",
    "            self.n_words += 1\n",
    "        else:\n",
    "            self.word2count[word] += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pBwz8kK8oTjH"
   },
   "source": [
    "The files are all in Unicode, to simplify we will turn Unicode\n",
    "characters to ASCII, make everything lowercase, and trim most\n",
    "punctuation.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RiyZRpvFoTjH"
   },
   "outputs": [],
   "source": [
    "# Turn a Unicode string to plain ASCII, thanks to\n",
    "# https://stackoverflow.com/a/518232/2809427\n",
    "def unicodeToAscii(s):\n",
    "    return ''.join(\n",
    "        c for c in unicodedata.normalize('NFD', s)\n",
    "        if unicodedata.category(c) != 'Mn'\n",
    "    )\n",
    "\n",
    "# Lowercase, trim, and remove non-letter characters\n",
    "\n",
    "\n",
    "def normalizeString(s):\n",
    "    s = unicodeToAscii(s.lower().strip())\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    return s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R2iz97nYoTjI"
   },
   "source": [
    "To read the data file we will split the file into lines, and then split\n",
    "lines into pairs. The files are all English → Other Language, so if we\n",
    "want to translate from Other Language → English I added the ``reverse``\n",
    "flag to reverse the pairs.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s8gLtSLyoTjI"
   },
   "outputs": [],
   "source": [
    "def readLangs(lang1, lang2, reverse=False):\n",
    "    print(\"Reading lines...\")\n",
    "\n",
    "    # Read the file and split into lines\n",
    "    lines = open('data/%s-%s.txt' % (lang1, lang2), encoding='utf-8').\\\n",
    "        read().strip().split('\\n')\n",
    "\n",
    "    # Split every line into pairs and normalize\n",
    "    pairs = [[normalizeString(s) for s in l.split('\\t')] for l in lines]\n",
    "\n",
    "    # Reverse pairs, make Lang instances\n",
    "    if reverse:\n",
    "        pairs = [list(reversed(p)) for p in pairs]\n",
    "        input_lang = Lang(lang2)\n",
    "        output_lang = Lang(lang1)\n",
    "    else:\n",
    "        input_lang = Lang(lang1)\n",
    "        output_lang = Lang(lang2)\n",
    "\n",
    "    return input_lang, output_lang, pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "90NyCjp7oTjJ"
   },
   "source": [
    "Since there are a *lot* of example sentences and we want to train\n",
    "something quickly, we'll trim the data set to only relatively short and\n",
    "simple sentences. Here the maximum length is 10 words (that includes\n",
    "ending punctuation) and we're filtering to sentences that translate to\n",
    "the form \"I am\" or \"He is\" etc. (accounting for apostrophes replaced\n",
    "earlier).\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0hC5odnfoTjJ"
   },
   "outputs": [],
   "source": [
    "MAX_LENGTH = 10\n",
    "\n",
    "eng_prefixes = (\n",
    "    \"i am \", \"i m \",\n",
    "    \"he is\", \"he s \",\n",
    "    \"she is\", \"she s \",\n",
    "    \"you are\", \"you re \",\n",
    "    \"we are\", \"we re \",\n",
    "    \"they are\", \"they re \"\n",
    ")\n",
    "\n",
    "\n",
    "def filterPair(p):\n",
    "    return len(p[0].split(' ')) < MAX_LENGTH and \\\n",
    "        len(p[1].split(' ')) < MAX_LENGTH and \\\n",
    "        p[1].startswith(eng_prefixes)\n",
    "\n",
    "\n",
    "def filterPairs(pairs):\n",
    "    return [pair for pair in pairs if filterPair(pair)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uFJFSQrvoTjK"
   },
   "source": [
    "The full process for preparing the data is:\n",
    "\n",
    "-  Read text file and split into lines, split lines into pairs\n",
    "-  Normalize text, filter by length and content\n",
    "-  Make word lists from sentences in pairs\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1zbqOWRVoTjK"
   },
   "outputs": [],
   "source": [
    "def prepareData(lang1, lang2, reverse=False):\n",
    "    input_lang, output_lang, pairs = readLangs(lang1, lang2, reverse)\n",
    "    print(\"Read %s sentence pairs\" % len(pairs))\n",
    "    pairs = filterPairs(pairs)\n",
    "    print(\"Trimmed to %s sentence pairs\" % len(pairs))\n",
    "    print(\"Counting words...\")\n",
    "    for pair in pairs:\n",
    "        input_lang.addSentence(pair[0])\n",
    "        output_lang.addSentence(pair[1])\n",
    "    print(\"Counted words:\")\n",
    "    print(input_lang.name, input_lang.n_words)\n",
    "    print(output_lang.name, output_lang.n_words)\n",
    "    return input_lang, output_lang, pairs\n",
    "\n",
    "\n",
    "input_lang, output_lang, pairs = prepareData('eng', 'fra', True)\n",
    "print(random.choice(pairs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KElEyvZOoTjL"
   },
   "source": [
    "### The Seq2Seq Model\n",
    "\n",
    "A Recurrent Neural Network, or RNN, is a network that operates on a\n",
    "sequence and uses its own output as input for subsequent steps.\n",
    "\n",
    "A `Sequence to Sequence network <https://arxiv.org/abs/1409.3215>`__, or\n",
    "seq2seq network, or `Encoder Decoder\n",
    "network <https://arxiv.org/pdf/1406.1078v3.pdf>`__, is a model\n",
    "consisting of two RNNs called the encoder and decoder. The encoder reads\n",
    "an input sequence and outputs a single vector, and the decoder reads\n",
    "that vector to produce an output sequence.\n",
    "\n",
    ".. figure:: /_static/img/seq-seq-images/seq2seq.png\n",
    "   :alt:\n",
    "\n",
    "Unlike sequence prediction with a single RNN, where every input\n",
    "corresponds to an output, the seq2seq model frees us from sequence\n",
    "length and order, which makes it ideal for translation between two\n",
    "languages.\n",
    "\n",
    "Consider the sentence \"Je ne suis pas le chat noir\" → \"I am not the\n",
    "black cat\". Most of the words in the input sentence have a direct\n",
    "translation in the output sentence, but are in slightly different\n",
    "orders, e.g. \"chat noir\" and \"black cat\". Because of the \"ne/pas\"\n",
    "construction there is also one more word in the input sentence. It would\n",
    "be difficult to produce a correct translation directly from the sequence\n",
    "of input words.\n",
    "\n",
    "With a seq2seq model the encoder creates a single vector which, in the\n",
    "ideal case, encodes the \"meaning\" of the input sequence into a single\n",
    "vector — a single point in some N dimensional space of sentences.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pRgxRkttoTjM"
   },
   "source": [
    "### The Encoder\n",
    "\n",
    "The encoder of a seq2seq network is a RNN that outputs some value for\n",
    "every word from the input sentence. For every input word the encoder\n",
    "outputs a vector and a hidden state, and uses the hidden state for the\n",
    "next input word.\n",
    "\n",
    ".. figure:: /_static/img/seq-seq-images/encoder-network.png\n",
    "   :alt:\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tcpRP9DKoTjM"
   },
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size)\n",
    "\n",
    "    def forward(self, input, hidden):\n",
    "        embedded = self.embedding(input).view(1, 1, -1)\n",
    "        output = embedded\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        return output, hidden\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "koRIOArfoTjM"
   },
   "source": [
    "### The Decoder\n",
    "\n",
    "The decoder is another RNN that takes the encoder output vector(s) and\n",
    "outputs a sequence of words to create the translation.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lUT9JS-poTjN"
   },
   "source": [
    "\n",
    "\n",
    "In the simplest seq2seq decoder we use only last output of the encoder.\n",
    "This last output is sometimes called the *context vector* as it encodes\n",
    "context from the entire sequence. This context vector is used as the\n",
    "initial hidden state of the decoder.\n",
    "\n",
    "At every step of decoding, the decoder is given an input token and\n",
    "hidden state. The initial input token is the start-of-string ``<SOS>``\n",
    "token, and the first hidden state is the context vector (the encoder's\n",
    "last hidden state).\n",
    "\n",
    ".. figure:: /_static/img/seq-seq-images/decoder-network.png\n",
    "   :alt:\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eJoM8Mn4oTjN"
   },
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.embedding = nn.Embedding(output_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "\n",
    "    def forward(self, input, hidden):\n",
    "        output = self.embedding(input).view(1, 1, -1)\n",
    "        output = F.relu(output)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        output = self.softmax(self.out(output[0]))\n",
    "        return output, hidden\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YLkleT6hoTjO"
   },
   "source": [
    "I encourage you to train and observe the results of this model, but to\n",
    "save space we'll be going straight for the gold and introducing the\n",
    "Attention Mechanism.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LCR7DceXoTjO"
   },
   "source": [
    "Attention Decoder\n",
    "^^^^^^^^^^^^^^^^^\n",
    "\n",
    "If only the context vector is passed betweeen the encoder and decoder,\n",
    "that single vector carries the burden of encoding the entire sentence.\n",
    "\n",
    "Attention allows the decoder network to \"focus\" on a different part of\n",
    "the encoder's outputs for every step of the decoder's own outputs. First\n",
    "we calculate a set of *attention weights*. These will be multiplied by\n",
    "the encoder output vectors to create a weighted combination. The result\n",
    "(called ``attn_applied`` in the code) should contain information about\n",
    "that specific part of the input sequence, and thus help the decoder\n",
    "choose the right output words.\n",
    "\n",
    ".. figure:: https://i.imgur.com/1152PYf.png\n",
    "   :alt:\n",
    "\n",
    "Calculating the attention weights is done with another feed-forward\n",
    "layer ``attn``, using the decoder's input and hidden state as inputs.\n",
    "Because there are sentences of all sizes in the training data, to\n",
    "actually create and train this layer we have to choose a maximum\n",
    "sentence length (input length, for encoder outputs) that it can apply\n",
    "to. Sentences of the maximum length will use all the attention weights,\n",
    "while shorter sentences will only use the first few.\n",
    "\n",
    ".. figure:: /_static/img/seq-seq-images/attention-decoder-network.png\n",
    "   :alt:\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n-tZt8K8oTjO"
   },
   "outputs": [],
   "source": [
    "class AttnDecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size, dropout_p=0.1, max_length=MAX_LENGTH):\n",
    "        super(AttnDecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.dropout_p = dropout_p\n",
    "        self.max_length = max_length\n",
    "\n",
    "        self.embedding = nn.Embedding(self.output_size, self.hidden_size)\n",
    "        self.attn = nn.Linear(self.hidden_size * 2, self.max_length)\n",
    "        self.attn_combine = nn.Linear(self.hidden_size * 2, self.hidden_size)\n",
    "        self.dropout = nn.Dropout(self.dropout_p)\n",
    "        self.gru = nn.GRU(self.hidden_size, self.hidden_size)\n",
    "        self.out = nn.Linear(self.hidden_size, self.output_size)\n",
    "\n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "        embedded = self.embedding(input).view(1, 1, -1)\n",
    "        embedded = self.dropout(embedded)\n",
    "\n",
    "        attn_weights = F.softmax(\n",
    "            self.attn(torch.cat((embedded[0], hidden[0]), 1)), dim=1)\n",
    "        attn_applied = torch.bmm(attn_weights.unsqueeze(0),\n",
    "                                 encoder_outputs.unsqueeze(0))\n",
    "\n",
    "        output = torch.cat((embedded[0], attn_applied[0]), 1)\n",
    "        output = self.attn_combine(output).unsqueeze(0)\n",
    "\n",
    "        output = F.relu(output)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "\n",
    "        output = F.log_softmax(self.out(output[0]), dim=1)\n",
    "        return output, hidden, attn_weights\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DrQR6uRVoTjP"
   },
   "source": [
    "\n",
    "\n",
    "### Training\n",
    "\n",
    "Preparing Training Data\n",
    "\n",
    "To train, for each pair we will need an input tensor (indexes of the\n",
    "words in the input sentence) and target tensor (indexes of the words in\n",
    "the target sentence). While creating these vectors we will append the\n",
    "EOS token to both sequences.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jPDyE3u6oTjP"
   },
   "outputs": [],
   "source": [
    "def indexesFromSentence(lang, sentence):\n",
    "    return [lang.word2index[word] for word in sentence.split(' ')]\n",
    "\n",
    "\n",
    "def tensorFromSentence(lang, sentence):\n",
    "    indexes = indexesFromSentence(lang, sentence)\n",
    "    indexes.append(EOS_token)\n",
    "    return torch.tensor(indexes, dtype=torch.long, device=device).view(-1, 1)\n",
    "\n",
    "\n",
    "def tensorsFromPair(pair):\n",
    "    input_tensor = tensorFromSentence(input_lang, pair[0])\n",
    "    target_tensor = tensorFromSentence(output_lang, pair[1])\n",
    "    return (input_tensor, target_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qq6IR_LgoTjP"
   },
   "source": [
    "### Training the Model\n",
    "\n",
    "To train we run the input sentence through the encoder, and keep track\n",
    "of every output and the latest hidden state. Then the decoder is given\n",
    "the ``<SOS>`` token as its first input, and the last hidden state of the\n",
    "encoder as its first hidden state.\n",
    "\n",
    "\"Teacher forcing\" is the concept of using the real target outputs as\n",
    "each next input, instead of using the decoder's guess as the next input.\n",
    "Using teacher forcing causes it to converge faster but `when the trained\n",
    "network is exploited, it may exhibit\n",
    "instability <http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.378.4095&rep=rep1&type=pdf>`__.\n",
    "\n",
    "You can observe outputs of teacher-forced networks that read with\n",
    "coherent grammar but wander far from the correct translation -\n",
    "intuitively it has learned to represent the output grammar and can \"pick\n",
    "up\" the meaning once the teacher tells it the first few words, but it\n",
    "has not properly learned how to create the sentence from the translation\n",
    "in the first place.\n",
    "\n",
    "Because of the freedom PyTorch's autograd gives us, we can randomly\n",
    "choose to use teacher forcing or not with a simple if statement. Turn\n",
    "``teacher_forcing_ratio`` up to use more of it.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8X76FoDFoTjP"
   },
   "outputs": [],
   "source": [
    "teacher_forcing_ratio = 0.5\n",
    "\n",
    "\n",
    "def train(input_tensor, target_tensor, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, max_length=MAX_LENGTH):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "\n",
    "    input_length = input_tensor.size(0)\n",
    "    target_length = target_tensor.size(0)\n",
    "\n",
    "    encoder_outputs = torch.zeros(max_length, encoder.hidden_size, device=device)\n",
    "\n",
    "    loss = 0\n",
    "\n",
    "    for ei in range(input_length):\n",
    "        encoder_output, encoder_hidden = encoder(\n",
    "            input_tensor[ei], encoder_hidden)\n",
    "        encoder_outputs[ei] = encoder_output[0, 0]\n",
    "\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "\n",
    "    decoder_hidden = encoder_hidden\n",
    "\n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    loss.backward()\n",
    "\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NhBofIiGoTjQ"
   },
   "source": [
    "This is a helper function to print time elapsed and estimated time\n",
    "remaining given the current time and progress %.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2cbdMCvNoTjQ"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import math\n",
    "\n",
    "\n",
    "def asMinutes(s):\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    return '%dm %ds' % (m, s)\n",
    "\n",
    "\n",
    "def timeSince(since, percent):\n",
    "    now = time.time()\n",
    "    s = now - since\n",
    "    es = s / (percent)\n",
    "    rs = es - s\n",
    "    return '%s (- %s)' % (asMinutes(s), asMinutes(rs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mklP0f8VoTjQ"
   },
   "source": [
    "The whole training process looks like this:\n",
    "\n",
    "-  Start a timer\n",
    "-  Initialize optimizers and criterion\n",
    "-  Create set of training pairs\n",
    "-  Start empty losses array for plotting\n",
    "\n",
    "Then we call ``train`` many times and occasionally print the progress (%\n",
    "of examples, time so far, estimated time) and average loss.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFA8uyDBoTjQ"
   },
   "outputs": [],
   "source": [
    "def trainIters(encoder, decoder, n_iters, print_every=1000, plot_every=100, learning_rate=0.01):\n",
    "    start = time.time()\n",
    "    plot_losses = []\n",
    "    print_loss_total = 0  # Reset every print_every\n",
    "    plot_loss_total = 0  # Reset every plot_every\n",
    "\n",
    "    encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "    decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "    training_pairs = [tensorsFromPair(random.choice(pairs))\n",
    "                      for i in range(n_iters)]\n",
    "    criterion = nn.NLLLoss()\n",
    "\n",
    "    for iter in range(1, n_iters + 1):\n",
    "        training_pair = training_pairs[iter - 1]\n",
    "        input_tensor = training_pair[0]\n",
    "        target_tensor = training_pair[1]\n",
    "\n",
    "        loss = train(input_tensor, target_tensor, encoder,\n",
    "                     decoder, encoder_optimizer, decoder_optimizer, criterion)\n",
    "        print_loss_total += loss\n",
    "        plot_loss_total += loss\n",
    "\n",
    "        if iter % print_every == 0:\n",
    "            print_loss_avg = print_loss_total / print_every\n",
    "            print_loss_total = 0\n",
    "            print('%s (%d %d%%) %.4f' % (timeSince(start, iter / n_iters),\n",
    "                                         iter, iter / n_iters * 100, print_loss_avg))\n",
    "\n",
    "        if iter % plot_every == 0:\n",
    "            plot_loss_avg = plot_loss_total / plot_every\n",
    "            plot_losses.append(plot_loss_avg)\n",
    "            plot_loss_total = 0\n",
    "\n",
    "    showPlot(plot_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "slMtEO_FoTjR"
   },
   "source": [
    "### Plotting results\n",
    "\n",
    "Plotting is done with matplotlib, using the array of loss values\n",
    "``plot_losses`` saved while training.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D-hQtgY9oTjR"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.switch_backend('agg')\n",
    "import matplotlib.ticker as ticker\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def showPlot(points):\n",
    "    plt.figure()\n",
    "    fig, ax = plt.subplots()\n",
    "    # this locator puts ticks at regular intervals\n",
    "    loc = ticker.MultipleLocator(base=0.2)\n",
    "    ax.yaxis.set_major_locator(loc)\n",
    "    plt.plot(points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r8wE-4jioTjR"
   },
   "source": [
    "### Evaluation\n",
    "\n",
    "Evaluation is mostly the same as training, but there are no targets so\n",
    "we simply feed the decoder's predictions back to itself for each step.\n",
    "Every time it predicts a word we add it to the output string, and if it\n",
    "predicts the EOS token we stop there. We also store the decoder's\n",
    "attention outputs for display later.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E_khje6GoTjR"
   },
   "outputs": [],
   "source": [
    "def evaluate(encoder, decoder, sentence, max_length=MAX_LENGTH):\n",
    "    with torch.no_grad():\n",
    "        input_tensor = tensorFromSentence(input_lang, sentence)\n",
    "        input_length = input_tensor.size()[0]\n",
    "        encoder_hidden = encoder.initHidden()\n",
    "\n",
    "        encoder_outputs = torch.zeros(max_length, encoder.hidden_size, device=device)\n",
    "\n",
    "        for ei in range(input_length):\n",
    "            encoder_output, encoder_hidden = encoder(input_tensor[ei],\n",
    "                                                     encoder_hidden)\n",
    "            encoder_outputs[ei] += encoder_output[0, 0]\n",
    "\n",
    "        decoder_input = torch.tensor([[SOS_token]], device=device)  # SOS\n",
    "\n",
    "        decoder_hidden = encoder_hidden\n",
    "\n",
    "        decoded_words = []\n",
    "        decoder_attentions = torch.zeros(max_length, max_length)\n",
    "\n",
    "        for di in range(max_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            decoder_attentions[di] = decoder_attention.data\n",
    "            topv, topi = decoder_output.data.topk(1)\n",
    "            if topi.item() == EOS_token:\n",
    "                decoded_words.append('<EOS>')\n",
    "                break\n",
    "            else:\n",
    "                decoded_words.append(output_lang.index2word[topi.item()])\n",
    "\n",
    "            decoder_input = topi.squeeze().detach()\n",
    "\n",
    "        return decoded_words, decoder_attentions[:di + 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IV5vpawIoTjR"
   },
   "source": [
    "We can evaluate random sentences from the training set and print out the\n",
    "input, target, and output to make some subjective quality judgements:\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pLDbHlxmoTjS"
   },
   "outputs": [],
   "source": [
    "def evaluateRandomly(encoder, decoder, n=10):\n",
    "    for i in range(n):\n",
    "        pair = random.choice(pairs)\n",
    "        print('>', pair[0])\n",
    "        print('=', pair[1])\n",
    "        output_words, attentions = evaluate(encoder, decoder, pair[0])\n",
    "        output_sentence = ' '.join(output_words)\n",
    "        print('<', output_sentence)\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-u_pJpTKoTjS"
   },
   "source": [
    "### Training and Evaluating\n",
    "\n",
    "With all these helper functions in place (it looks like extra work, but\n",
    "it makes it easier to run multiple experiments) we can actually\n",
    "initialize a network and start training.\n",
    "\n",
    "Remember that the input sentences were heavily filtered. For this small\n",
    "dataset we can use relatively small networks of 256 hidden nodes and a\n",
    "single GRU layer. After about 40 minutes on a MacBook CPU we'll get some\n",
    "reasonable results.\n",
    "\n",
    ".. Note::\n",
    "   If you run this notebook you can train, interrupt the kernel,\n",
    "   evaluate, and continue training later. Comment out the lines where the\n",
    "   encoder and decoder are initialized and run ``trainIters`` again.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I9HTlKHcoTjS"
   },
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "encoder1 = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "attn_decoder1 = AttnDecoderRNN(hidden_size, output_lang.n_words, dropout_p=0.1).to(device)\n",
    "\n",
    "trainIters(encoder1, attn_decoder1, 75000, print_every=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eUMOAFIHoTjS"
   },
   "outputs": [],
   "source": [
    "evaluateRandomly(encoder1, attn_decoder1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4TLHkNqzoTjS"
   },
   "source": [
    "### Visualizing Attention\n",
    "\n",
    "\n",
    "A useful property of the attention mechanism is its highly interpretable\n",
    "outputs. Because it is used to weight specific encoder outputs of the\n",
    "input sequence, we can imagine looking where the network is focused most\n",
    "at each time step.\n",
    "\n",
    "You could simply run ``plt.matshow(attentions)`` to see attention output\n",
    "displayed as a matrix, with the columns being input steps and rows being\n",
    "output steps:\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lGhWjlvXoTjT"
   },
   "outputs": [],
   "source": [
    "output_words, attentions = evaluate(\n",
    "    encoder1, attn_decoder1, \"je suis trop froid .\")\n",
    "plt.matshow(attentions.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dzJbc42eoTjT"
   },
   "source": [
    "For a better viewing experience we will do the extra work of adding axes\n",
    "and labels:\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xZlEiTYdoTjT"
   },
   "outputs": [],
   "source": [
    "def showAttention(input_sentence, output_words, attentions):\n",
    "    # Set up figure with colorbar\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    cax = ax.matshow(attentions.numpy(), cmap='bone')\n",
    "    fig.colorbar(cax)\n",
    "\n",
    "    # Set up axes\n",
    "    ax.set_xticklabels([''] + input_sentence.split(' ') +\n",
    "                       ['<EOS>'], rotation=90)\n",
    "    ax.set_yticklabels([''] + output_words)\n",
    "\n",
    "    # Show label at every tick\n",
    "    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def evaluateAndShowAttention(input_sentence):\n",
    "    output_words, attentions = evaluate(\n",
    "        encoder1, attn_decoder1, input_sentence)\n",
    "    print('input =', input_sentence)\n",
    "    print('output =', ' '.join(output_words))\n",
    "    showAttention(input_sentence, output_words, attentions)\n",
    "\n",
    "\n",
    "evaluateAndShowAttention(\"elle a cinq ans de moins que moi .\")\n",
    "\n",
    "evaluateAndShowAttention(\"elle est trop petit .\")\n",
    "\n",
    "evaluateAndShowAttention(\"je ne crains pas de mourir .\")\n",
    "\n",
    "evaluateAndShowAttention(\"c est un jeune directeur plein de talent .\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fOXhlS5ANBDx"
   },
   "source": [
    "# Transformers\n",
    "\n",
    "Transformers are a neural architecture comprised of encoder component and a decoder component, which use self-attention to reach high levels of performance for sequence to sequence tasks (and drops all RNN components!)\n",
    "\n",
    "**We highly recommend reading this [visual blog](http://jalammar.github.io/illustrated-transformer/) on the Transformer before proceeding.**\n",
    "\n",
    "\n",
    "## Models\n",
    "\n",
    "We usually see two families of models - the Generative Pre-Trained Transformer models (GPT-n) developed by Open-AI, and the BERT family of models developed by Google, which featured Bi-directional Transformers as language models.\n",
    "\n",
    "### GPT and similar models\n",
    "\n",
    "These are auto-regressive language models created by [OpenAI](https://openai.com/). GPT3 has 175 billion parameters and is one of the largest, most powerful generative language model available today.\n",
    "\n",
    "**We highly recommend reading this [visual blog](http://jalammar.github.io/illustrated-gpt2/) on the GPT-2 before proceeding.**\n",
    "\n",
    "\n",
    "### BERT and similar models\n",
    "\n",
    "The BERT model, first presented by Google Research, uses bi-directional transformers and two tasks - Masked Language Modelling task, and Next Sentence Prediction task, to train the model. These large transformers based models are trained on lots of data (Wikipedia + books). BERT has since spun off many models, include language specific ones (SpanBERT), domain specific models (SciBERT), and different level of sequences (CharacterBERT). \n",
    "\n",
    "**We highly recommend reading this [visual blog](http://jalammar.github.io/illustrated-bert/) on BERT and similar models before proceeding**.\n",
    "\n",
    "Since BERT, the state of the art has constatly been beat by models such as XL-Net. We recommend this [blog post](https://medium.com/@phylypo/a-survey-of-the-state-of-the-art-language-models-up-to-early-2020-aba824302c6) which describes various language models as of 2020 (right before GPT-3 burst on the scene). \n",
    "\n",
    "We also highly recommend these papers for a more critical view on such language models:\n",
    "\n",
    "[On the Dangers of Stochastic Parrots: Can Language Models Be Too Big? 🦜](https://dl.acm.org/doi/10.1145/3442188.3445922) - Bender et al, 2021\n",
    "\n",
    "[Climbing towards NLU: On Meaning, Form, and Understanding in the Age of Data](https://www.aclweb.org/anthology/2020.acl-main.463/) - Bender and Koller, 2020\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### On with BERT!\n",
    "\n",
    "So while Neural Networks can do a good job with some kind of classification tasks, they don't perform too well on intent classification. Let us see how a bidirectional transformer embedding like BERT might do. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert all of our data into torch tensors, the required datatype for our model\n",
    "\n",
    "train_inputs = torch.tensor(train_inputs)\n",
    "validation_inputs = torch.tensor(validation_inputs)\n",
    "train_labels = torch.tensor(train_labels)\n",
    "validation_labels = torch.tensor(validation_labels)\n",
    "train_masks = torch.tensor(train_masks)\n",
    "validation_masks = torch.tensor(validation_masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select a batch size for training. For fine-tuning BERT on a specific task, the authors recommend a batch size of 16 or 32\n",
    "batch_size = 32\n",
    "\n",
    "# Create an iterator of our data with torch DataLoader. This helps save on memory during training because, unlike a for loop, \n",
    "# with an iterator the entire dataset does not need to be loaded into memory\n",
    "\n",
    "train_data = TensorDataset(train_inputs, train_masks, train_labels)\n",
    "train_sampler = RandomSampler(train_data)\n",
    "train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)\n",
    "\n",
    "validation_data = TensorDataset(validation_inputs, validation_masks, validation_labels)\n",
    "validation_sampler = SequentialSampler(validation_data)\n",
    "validation_dataloader = DataLoader(validation_data, sampler=validation_sampler, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading our Models\n",
    "\n",
    "### Train Model\n",
    "Now that our input data is properly formatted, it's time to fine tune the BERT model.\n",
    "For this task, we first want to modify the pre-trained BERT model to give outputs for classification, and then we want to continue training the model on our dataset until the entire model, end-to-end, is well-suited for our task. Thankfully, the huggingface pytorch implementation includes a set of interfaces designed for a variety of NLP tasks. Though these interfaces are all built on top of a trained BERT model, each has different top layers and output types designed to accomodate their specific NLP task.\n",
    "We'll load BertForSequenceClassification. This is the normal BERT model with an added single linear layer on top for classification that we will use as a sentence classifier. As we feed input data, the entire pre-trained BERT model and the additional untrained classification layer is trained on our specific task.\n",
    "\n",
    "### Structure of Fine-Tuning Model\n",
    "\n",
    "As we've showed beforehand, the first token of every sequence is the special classification token ([CLS]). Unlike the hidden state vector corresponding to a normal word token, the hidden state corresponding to this special token is designated by the authors of BERT as an aggregate representation of the whole sentence used for classification tasks. As such, when we feed in an input sentence to our model during training, the output is the length 768 hidden state vector corresponding to this token. The additional layer that we've added on top consists of untrained linear neurons of size [hidden_state, number_of_labels], so [768,2], meaning that the output of BERT plus our classification layer is a vector of two numbers representing the \"score\" for \"grammatical/non-grammatical\" that are then fed into cross-entropy loss.\n",
    "\n",
    "### The Fine-Tuning Process\n",
    "\n",
    "Because the pre-trained BERT layers already encode a lot of information about the language, training the classifier is relatively inexpensive. Rather than training every layer in a large model from scratch, it's as if we have already trained the bottom layers 95% of where they need to be, and only really need to train the top layer, with a bit of tweaking going on in the lower levels to accomodate our task.\n",
    "Sometimes practicioners will opt to \"freeze\" certain layers when fine-tuning, or to apply different learning rates, apply diminishing learning rates, etc. all in an effort to preserve the good quality weights in the network and speed up training (often considerably). In fact, recent research on BERT specifically has demonstrated that freezing the majority of the weights results in only minimal accuracy declines, but there are exceptions and broader rules of transfer learning that should also be considered. For example, if your task and fine-tuning dataset is very different from the dataset used to train the transfer learning model, freezing the weights may not be a good idea. OK, let's load BERT! There are a few different pre-trained BERT models available. \"bert-base-uncased\" means the version that has only lowercase letters (\"uncased\") and is the smaller version of the two (\"base\" vs \"large\").\n",
    "\n",
    "Credit to Michel Kana's [tutorial](https://towardsdatascience.com/bert-for-dummies-step-by-step-tutorial-fb90890ffe03) and the [tutorial](https://colab.research.google.com/drive/1Y4o3jh3ZH70tl6mCd76vz_IxX23biCPP#scrollTo=GuE5BqICAne2) by Chris McCormick and Nick Ryan who describe the workings of BERT and the way it is used by the ```transformers``` package. \n",
    "\n",
    "## WARNING: SHIFT TO A GPU ENABLED MACHINE (e.g., Google Colab)\n",
    "\n",
    "Note that you only want to run the following code if you have a GPU. Otherwise, rerun the **same** cells we just ran on the Colab file to train your model, download it to your local, and load it by running\n",
    "```model = BertForSequenceClassification.from_pretrained(\"my_model_directory\", num_labels=2)```.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BertForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=2)\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_optimizer = list(model.named_parameters())\n",
    "no_decay = ['bias', 'gamma', 'beta']\n",
    "optimizer_grouped_parameters = [\n",
    "    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],\n",
    "     'weight_decay_rate': 0.01},\n",
    "    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],\n",
    "     'weight_decay_rate': 0.0}\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This variable contains all of the hyperparemeter information our training loop needs\n",
    "optimizer = AdamW(model.parameters(),\n",
    "                  lr = 2e-5, # args.learning_rate - default is 5e-5, our notebook had 2e-5\n",
    "                  eps = 1e-8 # args.adam_epsilon  - default is 1e-8.\n",
    "                )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import get_linear_schedule_with_warmup\n",
    "\n",
    "# Number of training epochs (authors recommend between 2 and 4)\n",
    "epochs = 4\n",
    "\n",
    "# Total number of training steps is number of batches * number of epochs.\n",
    "total_steps = len(train_dataloader) * epochs\n",
    "\n",
    "# Create the learning rate scheduler.\n",
    "scheduler = get_linear_schedule_with_warmup(optimizer, \n",
    "                                            num_warmup_steps = 0, # Default value in run_glue.py\n",
    "                                            num_training_steps = total_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to calculate the accuracy of our predictions vs labels\n",
    "def flat_accuracy(preds, labels):\n",
    "    pred_flat = np.argmax(preds, axis=1).flatten()\n",
    "    labels_flat = labels.flatten()\n",
    "    return np.sum(pred_flat == labels_flat) / len(labels_flat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import datetime\n",
    "\n",
    "def format_time(elapsed):\n",
    "    '''\n",
    "    Takes a time in seconds and returns a string hh:mm:ss\n",
    "    '''\n",
    "    # Round to the nearest second.\n",
    "    elapsed_rounded = int(round((elapsed)))\n",
    "    \n",
    "    # Format as hh:mm:ss\n",
    "    return str(datetime.timedelta(seconds=elapsed_rounded))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the following cell can take 12 hours or longer if run without a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "# This training code is based on the `run_glue.py` script here:\n",
    "# https://github.com/huggingface/transformers/blob/5bfcd0485ece086ebcbed2d008813037968a9e58/examples/run_glue.py#L128\n",
    "\n",
    "# Set the seed value all over the place to make this reproducible.\n",
    "seed_val = 42\n",
    "\n",
    "random.seed(seed_val)\n",
    "np.random.seed(seed_val)\n",
    "torch.manual_seed(seed_val)\n",
    "torch.cuda.manual_seed_all(seed_val)\n",
    "\n",
    "# Store the average loss after each epoch so we can plot them.\n",
    "loss_values = []\n",
    "\n",
    "# For each epoch...\n",
    "for epoch_i in range(0, epochs):\n",
    "    \n",
    "    # ========================================\n",
    "    #               Training\n",
    "    # ========================================\n",
    "    \n",
    "    # Perform one full pass over the training set.\n",
    "\n",
    "    print(\"\")\n",
    "    print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))\n",
    "    print('Training...')\n",
    "\n",
    "    # Measure how long the training epoch takes.\n",
    "    t0 = time.time()\n",
    "\n",
    "    # Reset the total loss for this epoch.\n",
    "    total_loss = 0\n",
    "\n",
    "    # Put the model into training mode. Don't be mislead--the call to \n",
    "    # `train` just changes the *mode*, it doesn't *perform* the training.\n",
    "    # `dropout` and `batchnorm` layers behave differently during training\n",
    "    # vs. test (source: https://stackoverflow.com/questions/51433378/what-does-model-train-do-in-pytorch)\n",
    "    model.train()\n",
    "\n",
    "    # For each batch of training data...\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "\n",
    "        # Progress update every 40 batches.\n",
    "        if step % 40 == 0 and not step == 0:\n",
    "            # Calculate elapsed time in minutes.\n",
    "            elapsed = format_time(time.time() - t0)\n",
    "            \n",
    "            # Report progress.\n",
    "            print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_dataloader), elapsed))\n",
    "\n",
    "        # Unpack this training batch from our dataloader. \n",
    "        #\n",
    "        # As we unpack the batch, we'll also copy each tensor to the GPU using the \n",
    "        # `to` method.\n",
    "        #\n",
    "        # `batch` contains three pytorch tensors:\n",
    "        #   [0]: input ids \n",
    "        #   [1]: attention masks\n",
    "        #   [2]: labels \n",
    "        b_input_ids = batch[0].to(device)\n",
    "        b_input_mask = batch[1].to(device)\n",
    "        b_labels = batch[2].to(device)\n",
    "\n",
    "        # Always clear any previously calculated gradients before performing a\n",
    "        # backward pass. PyTorch doesn't do this automatically because \n",
    "        # accumulating the gradients is \"convenient while training RNNs\". \n",
    "        # (source: https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch)\n",
    "        model.zero_grad()        \n",
    "\n",
    "        # Perform a forward pass (evaluate the model on this training batch).\n",
    "        # This will return the loss (rather than the model output) because we\n",
    "        # have provided the `labels`.\n",
    "        # The documentation for this `model` function is here: \n",
    "        # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification\n",
    "        outputs = model(b_input_ids, \n",
    "                    token_type_ids=None, \n",
    "                    attention_mask=b_input_mask, \n",
    "                    labels=b_labels)\n",
    "        \n",
    "        # The call to `model` always returns a tuple, so we need to pull the \n",
    "        # loss value out of the tuple.\n",
    "        loss = outputs[0]\n",
    "\n",
    "        # Accumulate the training loss over all of the batches so that we can\n",
    "        # calculate the average loss at the end. `loss` is a Tensor containing a\n",
    "        # single value; the `.item()` function just returns the Python value \n",
    "        # from the tensor.\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # Perform a backward pass to calculate the gradients.\n",
    "        loss.backward()\n",
    "\n",
    "        # Clip the norm of the gradients to 1.0.\n",
    "        # This is to help prevent the \"exploding gradients\" problem.\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "        # Update parameters and take a step using the computed gradient.\n",
    "        # The optimizer dictates the \"update rule\"--how the parameters are\n",
    "        # modified based on their gradients, the learning rate, etc.\n",
    "        optimizer.step()\n",
    "\n",
    "        # Update the learning rate.\n",
    "        scheduler.step()\n",
    "\n",
    "    # Calculate the average loss over the training data.\n",
    "    avg_train_loss = total_loss / len(train_dataloader)            \n",
    "    \n",
    "    # Store the loss value for plotting the learning curve.\n",
    "    loss_values.append(avg_train_loss)\n",
    "\n",
    "    print(\"\")\n",
    "    print(\"  Average training loss: {0:.2f}\".format(avg_train_loss))\n",
    "    print(\"  Training epcoh took: {:}\".format(format_time(time.time() - t0)))\n",
    "        \n",
    "    # ========================================\n",
    "    #               Validation\n",
    "    # ========================================\n",
    "    # After the completion of each training epoch, measure our performance on\n",
    "    # our validation set.\n",
    "\n",
    "    print(\"\")\n",
    "    print(\"Running Validation...\")\n",
    "\n",
    "    t0 = time.time()\n",
    "\n",
    "    # Put the model in evaluation mode--the dropout layers behave differently\n",
    "    # during evaluation.\n",
    "    model.eval()\n",
    "\n",
    "    # Tracking variables \n",
    "    eval_loss, eval_accuracy = 0, 0\n",
    "    nb_eval_steps, nb_eval_examples = 0, 0\n",
    "\n",
    "    # Evaluate data for one epoch\n",
    "    for batch in validation_dataloader:\n",
    "        \n",
    "        # Add batch to GPU\n",
    "        batch = tuple(t.to(device) for t in batch)\n",
    "        \n",
    "        # Unpack the inputs from our dataloader\n",
    "        b_input_ids, b_input_mask, b_labels = batch\n",
    "        \n",
    "        # Telling the model not to compute or store gradients, saving memory and\n",
    "        # speeding up validation\n",
    "        with torch.no_grad():        \n",
    "\n",
    "            # Forward pass, calculate logit predictions.\n",
    "            # This will return the logits rather than the loss because we have\n",
    "            # not provided labels.\n",
    "            # token_type_ids is the same as the \"segment ids\", which \n",
    "            # differentiates sentence 1 and 2 in 2-sentence tasks.\n",
    "            # The documentation for this `model` function is here: \n",
    "            # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification\n",
    "            outputs = model(b_input_ids, \n",
    "                            token_type_ids=None, \n",
    "                            attention_mask=b_input_mask)\n",
    "        \n",
    "        # Get the \"logits\" output by the model. The \"logits\" are the output\n",
    "        # values prior to applying an activation function like the softmax.\n",
    "        logits = outputs[0]\n",
    "\n",
    "        # Move logits and labels to CPU\n",
    "        logits = logits.detach().cpu().numpy()\n",
    "        label_ids = b_labels.to('cpu').numpy()\n",
    "        \n",
    "        # Calculate the accuracy for this batch of test sentences.\n",
    "        tmp_eval_accuracy = flat_accuracy(logits, label_ids)\n",
    "        \n",
    "        # Accumulate the total accuracy.\n",
    "        eval_accuracy += tmp_eval_accuracy\n",
    "\n",
    "        # Track the number of batches\n",
    "        nb_eval_steps += 1\n",
    "\n",
    "    # Report the final accuracy for this validation run.\n",
    "    print(\"  Accuracy: {0:.2f}\".format(eval_accuracy/nb_eval_steps))\n",
    "    print(\"  Validation took: {:}\".format(format_time(time.time() - t0)))\n",
    "\n",
    "print(\"\")\n",
    "print(\"Training complete!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "# Use plot styling from seaborn.\n",
    "sns.set(style='darkgrid')\n",
    "\n",
    "# Increase the plot size and font size.\n",
    "sns.set(font_scale=1.5)\n",
    "plt.rcParams[\"figure.figsize\"] = (12,6)\n",
    "\n",
    "# Plot the learning curve.\n",
    "plt.plot(loss_values, 'b-o')\n",
    "\n",
    "# Label the plot.\n",
    "plt.title(\"Training loss\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Loss\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## COME BACK TO THIS NOTEBOOK if you are using a non-GPU machine to load and work with your GPU-trained model\n",
    "\n",
    "Once you tune your model on Colab (or on your own machine if you decided to do that instead), you load it here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BertForSequenceClassification.from_pretrained(\"model_save\", num_labels=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Holdout Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Load the dataset into a pandas dataframe.\n",
    "df = pd.read_csv(\"../data/cola_public/raw/out_of_domain_dev.tsv\", delimiter='\\t', header=None, names=['sentence_source', 'label', 'label_notes', 'sentence'])\n",
    "\n",
    "# Report the number of sentences.\n",
    "print('Number of test sentences: {:,}\\n'.format(df.shape[0]))\n",
    "\n",
    "# Create sentence and label lists\n",
    "sentences = df.sentence.values\n",
    "labels = df.label.values\n",
    "\n",
    "# Tokenize all of the sentences and map the tokens to thier word IDs.\n",
    "input_ids = []\n",
    "\n",
    "# For every sentence...\n",
    "for sent in sentences:\n",
    "    # `encode` will:\n",
    "    #   (1) Tokenize the sentence.\n",
    "    #   (2) Prepend the `[CLS]` token to the start.\n",
    "    #   (3) Append the `[SEP]` token to the end.\n",
    "    #   (4) Map tokens to their IDs.\n",
    "    encoded_sent = tokenizer.encode(\n",
    "                        sent,                      # Sentence to encode.\n",
    "                        add_special_tokens = True, # Add '[CLS]' and '[SEP]'\n",
    "                   )\n",
    "    \n",
    "    input_ids.append(encoded_sent)\n",
    "\n",
    "# Pad our input tokens\n",
    "input_ids = pad_sequences(input_ids, maxlen=MAX_LEN, \n",
    "                          dtype=\"long\", truncating=\"post\", padding=\"post\")\n",
    "\n",
    "# Create attention masks\n",
    "attention_masks = []\n",
    "\n",
    "# Create a mask of 1s for each token followed by 0s for padding\n",
    "for seq in input_ids:\n",
    "    seq_mask = [float(i>0) for i in seq]\n",
    "    attention_masks.append(seq_mask) \n",
    "\n",
    "# Convert to tensors.\n",
    "prediction_inputs = torch.tensor(input_ids)\n",
    "prediction_masks = torch.tensor(attention_masks)\n",
    "prediction_labels = torch.tensor(labels)\n",
    "\n",
    "# Set the batch size.  \n",
    "batch_size = 32  \n",
    "\n",
    "# Create the DataLoader.\n",
    "prediction_data = TensorDataset(prediction_inputs, prediction_masks, prediction_labels)\n",
    "prediction_sampler = SequentialSampler(prediction_data)\n",
    "prediction_dataloader = DataLoader(prediction_data, sampler=prediction_sampler, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prediction on test set\n",
    "\n",
    "print('Predicting labels for {:,} test sentences...'.format(len(prediction_inputs)))\n",
    "\n",
    "# Put model in evaluation mode\n",
    "model.eval()\n",
    "\n",
    "# Tracking variables \n",
    "predictions , true_labels = [], []\n",
    "\n",
    "# Predict \n",
    "for batch in prediction_dataloader:\n",
    "    # Add batch to GPU\n",
    "    batch = tuple(t.to(device) for t in batch)\n",
    "\n",
    "    # Unpack the inputs from our dataloader\n",
    "    b_input_ids, b_input_mask, b_labels = batch\n",
    "\n",
    "    # Telling the model not to compute or store gradients, saving memory and \n",
    "    # speeding up prediction\n",
    "    with torch.no_grad():\n",
    "      # Forward pass, calculate logit predictions\n",
    "      outputs = model(b_input_ids, token_type_ids=None, \n",
    "                      attention_mask=b_input_mask)\n",
    "\n",
    "    logits = outputs[0]\n",
    "\n",
    "    # Move logits and labels to CPU\n",
    "    logits = logits.detach().cpu().numpy()\n",
    "    label_ids = b_labels.to('cpu').numpy()\n",
    "\n",
    "    # Store predictions and true labels\n",
    "    predictions.append(logits)\n",
    "    true_labels.append(label_ids)\n",
    "\n",
    "print('    DONE.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Positive samples: %d of %d (%.2f%%)' % (df.label.sum(), len(df.label), (df.label.sum() / len(df.label) * 100.0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import matthews_corrcoef\n",
    "\n",
    "matthews_set = []\n",
    "\n",
    "# Evaluate each test batch using Matthew's correlation coefficient\n",
    "print('Calculating Matthews Corr. Coef. for each batch...')\n",
    "\n",
    "# For each input batch...\n",
    "for i in range(len(true_labels)):\n",
    "\n",
    "    # The predictions for this batch are a 2-column ndarray (one column for \"0\" \n",
    "    # and one column for \"1\"). Pick the label with the highest value and turn this\n",
    "    # in to a list of 0s and 1s.\n",
    "    pred_labels_i = np.argmax(predictions[i], axis=1).flatten()\n",
    "\n",
    "    # Calculate and store the coef for this batch.  \n",
    "    matthews = matthews_corrcoef(true_labels[i], pred_labels_i)                \n",
    "    matthews_set.append(matthews)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matthews_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine the predictions for each batch into a single list of 0s and 1s.\n",
    "flat_predictions = [item for sublist in predictions for item in sublist]\n",
    "flat_predictions = np.argmax(flat_predictions, axis=1).flatten()\n",
    "\n",
    "# Combine the correct labels for each batch into a single list.\n",
    "flat_true_labels = [item for sublist in true_labels for item in sublist]\n",
    "\n",
    "# Calculate the MCC\n",
    "mcc = matthews_corrcoef(flat_true_labels, flat_predictions)\n",
    "\n",
    "print('MCC: %.3f' % mcc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's pretty good performance. Note that we used [Matthews Correlation](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html) to meausure this. It ranges from -1 to 1, with +1 being the best. The Google BERT model has a similar score too, so this model performed quite well. It took a long time though, approximately a day with no GPU. It would be significantly faster if a CUDA enabled machine ran this. Hence, we recommend that you run this on the Collab notebook.\n",
    "\n",
    "The following lines save the model to disk, if you would like to: note that we ran this in the colab file to save it to disk there as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "\n",
    "# # Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()\n",
    "\n",
    "# output_dir = './model_save/'\n",
    "\n",
    "# # Create output directory if needed\n",
    "# if not os.path.exists(output_dir):\n",
    "#     os.makedirs(output_dir)\n",
    "\n",
    "# print(\"Saving model to %s\" % output_dir)\n",
    "\n",
    "# # Save a trained model, configuration and tokenizer using `save_pretrained()`.\n",
    "# # They can then be reloaded using `from_pretrained()`\n",
    "# model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training\n",
    "# model_to_save.save_pretrained(output_dir)\n",
    "# tokenizer.save_pretrained(output_dir)\n",
    "\n",
    "# # Good practice: save your training arguments together with the trained model\n",
    "# # torch.save(args, os.path.join(output_dir, 'training_args.bin'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <font color=\"red\">*Exercise 1*</font>\n",
    "\n",
    "<font color=\"red\">Construct cells immediately below this that estimate a deep classification model with LSTM and BERT in order to predict pre-established data labels relevant to your final project. Which works better? Are the errors the same or different?\n",
    "\n",
    "<font color=\"red\">***Stretch***</font>: <font color=\"red\">Now alter the neural network by stacking network layers, adjusting the embedding dimension, compare its performance with your model above, and interpret why it might be different. Also consider whether you could use Seq2Seq in a similar way. Why or why not?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embeddings, Context Words\n",
    "\n",
    "We saw how a bootstrapped BERT model performed so much better than a model trained from scatch. Because BERT's method of capturing context is bidirectional, meaning that words can now have different word embedding values based on their location within a sentence. Let us use the same BERT model to capture sentence and word embeddings. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's go through the sentence format for the BERT model, as well as how our vocabulary looks like. Note that you have to use the BERT tokenizer to use the BERT model because of the similar vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Here is the sentence I want embeddings for.\"\n",
    "marked_text = \"[CLS] \" + text + \" [SEP]\"\n",
    "\n",
    "# Tokenize our sentence with the BERT tokenizer.\n",
    "tokenized_text = tokenizer.tokenize(marked_text)\n",
    "\n",
    "# Print out the tokens.\n",
    "print (tokenized_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "BERTS model uses a WordPiece technique to do its tokenizing, as described in the paper. That's why the word embedding is split up the way it is.\n",
    "A quick peek at what the voabulary looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(tokenizer.vocab.keys())[6000:6030]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a new example sentence with multiple meanings of the word \"bank\"\n",
    "text = \"After stealing money from the bank vault, the bank robber was seen fishing on the Mississippi river bank.\"\n",
    "\n",
    "# Add the special tokens.\n",
    "marked_text = \"[CLS] \" + text + \" [SEP]\"\n",
    "\n",
    "# Split the sentence into tokens.\n",
    "tokenized_text = tokenizer.tokenize(marked_text)\n",
    "\n",
    "# Map the token strings to their vocabulary indeces.\n",
    "indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "\n",
    "# Display the words with their indices.\n",
    "for tup in zip(tokenized_text, indexed_tokens):\n",
    "    print('{:<12} {:>6,}'.format(tup[0], tup[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Segment ID\n",
    "\n",
    "BERT is trained on and expects sentence pairs, using 1s and 0s to distinguish between the two sentences. That is, for each token in “tokenized_text,” we must specify which sentence it belongs to: sentence 0 (a series of 0s) or sentence 1 (a series of 1s). For our purposes, single-sentence inputs only require a series of 1s, so we will create a vector of 1s for each token in our input sentence.\n",
    "\n",
    "If you want to process two sentences, assign each word in the first sentence plus the ‘[SEP]’ token a 0, and all tokens of the second sentence a 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segments_ids = [1] * len(tokenized_text)\n",
    "\n",
    "print (segments_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like we did for classification, we now convert these segments to tensors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The embedding layer is the hidden state layer, and this is what we pick up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert inputs to PyTorch tensors\n",
    "tokens_tensor = torch.tensor([indexed_tokens])\n",
    "segments_tensors = torch.tensor([segments_ids])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load pre-trained model (weights)\n",
    "model_embedding = BertModel.from_pretrained('bert-base-uncased')\n",
    "\n",
    "# Put the model in \"evaluation\" mode, meaning feed-forward operation.\n",
    "model_embedding.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model_embedding(tokens_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(output[0][0][0]), len(output[1][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Understanding the Output\n",
    "\n",
    "This kind of forward pass returns us the last layer of the net, which we will use to make our vectors. The first object returned contains the batch number, followed by each of the tokens and their vector values. The second object contains a vector value, which I suspect is the sentence vector of the tokens. \n",
    "\n",
    "The first index is the batch size, and our batch size is 1, so we just choose the 0th index and work with that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_embeddings = output[0]\n",
    "sentence_embedding = output[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(word_embeddings[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_embeddings[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s take a quick look at the range of values for a given layer and token.\n",
    "\n",
    "You’ll find that the range is fairly similar for all layers and tokens, with the majority of values falling between [-2, 2], and a small smattering of values around -10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vec = word_embeddings[0][0]\n",
    "vec = vec.detach().numpy()\n",
    "# Plot the values as a histogram to show their distribution.\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.hist(vec, bins=200)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These values are grouped by layer - we can use the permute function to make it grouped by each individual token instead. Let us look at what the later looks like:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Word Vectors\n",
    "\n",
    "So each of those tokens have embedding values - let us try and compare them with each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_vecs = []\n",
    "# For each token in the sentence...\n",
    "for embedding in word_embeddings[0]:\n",
    "    cat_vec = embedding.detach().numpy()\n",
    "    # Use `cat_vec` to represent `token`.\n",
    "    token_vecs.append(cat_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(token_vecs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another method to create the vectors is to sum the last four layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sentence Vector\n",
    "\n",
    "To get a single vector for our entire sentence we have multiple application-dependent strategies - we could just average all the tokens in our sentence. We can also use this oppurtunity to see if the second vector returned is a sentence vector too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence_embedding_0 = sentence_embedding.detach().numpy()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence_embedding_1 = np.mean(token_vecs, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(sentence_embedding_0), len(sentence_embedding_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember that the power of these vectors is how they are context dependant - our sentence had multiple uses of the word bank. Let us see the index and the word of the sentence and check the context accordingly. We'll then print the simlarity values for the similar and different meanings and see how it turns out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, token_str in enumerate(tokenized_text):\n",
    "    print(i, token_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('First 5 vector values for each instance of \"bank\".')\n",
    "print('')\n",
    "print(\"bank vault   \", str(token_vecs[6][:5]))\n",
    "print(\"bank robber  \", str(token_vecs[10][:5]))\n",
    "print(\"river bank   \", str(token_vecs[19][:5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import cosine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the cosine similarity between the word bank \n",
    "# in \"bank robber\" vs \"river bank\" (different meanings).\n",
    "diff_bank = 1 - cosine(token_vecs[10], token_vecs[19])\n",
    "\n",
    "# Calculate the cosine similarity between the word bank\n",
    "# in \"bank robber\" vs \"bank vault\" (same meaning).\n",
    "same_bank = 1 - cosine(token_vecs[10], token_vecs[6])\n",
    "\n",
    "print('Vector similarity for  *similar*  meanings:  %.2f' % same_bank)\n",
    "print('Vector similarity for *different* meanings:  %.2f' % diff_bank)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This makes sense! Let us see if the mean value of all the tokens and what we think is the sentence vector is the same thing, by checking their cosine distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1 - cosine(sentence_embedding_0, sentence_embedding_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That is good - it seems it is indeed the sentence vector, so we can now write two functions which calculate the word and sentence vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def word_vector(text, word_id, model, tokenizer):\n",
    "    marked_text = \"[CLS] \" + text + \" [SEP]\"\n",
    "    tokenized_text = tokenizer.tokenize(marked_text)\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "    tokens_tensor = torch.tensor([indexed_tokens])\n",
    "    word_embeddings = model(tokens_tensor)\n",
    "    sentence_embeddings = model(tokens_tensor)[1]\n",
    "    vector = word_embeddings[0][word_id].detach().numpy()\n",
    "    return vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_10 = word_vector(text, 6, model_embedding, tokenizer)\n",
    "word_6 = word_vector(text, 10, model_embedding, tokenizer)\n",
    "word_19 = word_vector(text, 19, model_embedding, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sentence_vector(text, model, tokenizer, method=\"average\"):\n",
    "    marked_text = \"[CLS] \" + text + \" [SEP]\"\n",
    "    tokenized_text = tokenizer.tokenize(marked_text)\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "    tokens_tensor = torch.tensor([indexed_tokens])\n",
    "    word_embeddings = model(tokens_tensor)[0]\n",
    "    sentence_embeddings = model(tokens_tensor)[1]\n",
    "    token_vecs = []\n",
    "    \n",
    "    for embedding in word_embeddings[0]:\n",
    "        cat_vec = embedding.detach().numpy()\n",
    "        token_vecs.append(cat_vec)\n",
    "        \n",
    "    if method == \"average\":\n",
    "        sentence_embedding = np.mean(token_vecs, axis=0)\n",
    "    if method == \"model\":\n",
    "        sentence_embedding = sentence_embeddings\n",
    "    # do something\n",
    "    return sentence_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sen_vec_0 = sentence_vector(text, model_embedding, tokenizer)\n",
    "sen_vec_1 = sentence_vector(text, model_embedding, tokenizer, method=\"model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Similarity metrics\n",
    "It is worth noting that word-level similarity comparisons are not appropriate with BERT embeddings because these embeddings are contextually dependent, meaning that the word vector changes depending on the sentence it appears in. This enables direct sensitivity to polysemy so that, e.g., your representation encodes river “bank” and not a financial institution “bank”. Nevertheless, it makes direct word-to-word similarity comparisons less valuable. For sentence embeddings, however, similarity comparison is still valid such that one can query, for example, a single sentence against a dataset of other sentences in order to find the most similar. Depending on the similarity metric used, the resulting similarity values will be less informative than the relative ranking of similarity outputs as some similarity metrics make assumptions about the vector space (equally-weighted dimensions, for example) that do not hold for our 768-dimensional vector space.\n",
    "\n",
    "### Using the Vectors\n",
    "\n",
    "Without fine-tuning, BERT features may be less useful than plain GloVe or word2vec.\n",
    "They start to be interesting when you fine-tune a classifier on top of BERT. \n",
    "\n",
    "<!---\n",
    "### Using Transformers Pipelines\n",
    "\n",
    "The context vectors make the other pipeline functions which transformers has built in a lot more powerful. \n",
    "\n",
    "### NOTE\n",
    "The pipeline functionality in transformers is currently being worked on and might be broken, so it is an optional part of the exercise. Do try to uncomment the lines of code and try to see if it works, though! If you have managed to get transformers v2.5.1 installed, it will work - I managed to get it to work sometimes, it can be annoying to get it to work but when it works it works well.\n",
    "\n",
    "Consider the following:\n",
    "--->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hcBeb-E0GApS"
   },
   "source": [
    "### BERT Sentence Perplexity\n",
    "\n",
    "BERT can be used as a language model to see how \"likely\" a sentence is in that language model's configurations. This is the core technique used in one of our course readings, P. Vicinanza, A. Goldberg, S. Srivastava. 2020. “[Who Sees the Future? A Deep Learning Language Model  Demonstrates the Vision Advantage of Being Small](https://osf.io/preprints/socarxiv/j24pw/)”. \n",
    "\n",
    "Here is some additional reading on the topic. As the use of deep learning in social science grows, perplexity may be a particularly promising metric for high-quality research projects!\n",
    "\n",
    "- [Can We Use BERT as a Language Model to Assign a Score to a Sentence?](https://www.scribendi.ai/can-we-use-bert-as-a-language-model-to-assign-score-of-a-sentence/)\n",
    "\n",
    "- [Comparing BERT and GPT-2 as Language Models to Score the Grammatical Correctness of a Sentence](https://www.scribendi.ai/comparing-bert-and-gpt-2-as-language-models-to-score-the-grammatical-correctness-of-a-sentence/)\n",
    "\n",
    "- [Transformers - Perplexity of fixed-length models](https://huggingface.co/transformers/perplexity.html)\n",
    "\n",
    "- [Understanding evaluation metrics for language models](https://thegradient.pub/understanding-evaluation-metrics-for-language-models/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SYOx7a1VKE6r"
   },
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mkbud6iEJJma"
   },
   "outputs": [],
   "source": [
    "from transformers import BertTokenizer,BertForMaskedLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AKxudeIGJNML",
    "outputId": "d60494d3-fa6a-49fe-8d6a-a5b8005faeb4"
   },
   "outputs": [],
   "source": [
    "bertMaskedLM = BertForMaskedLM.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xR4c5pp3JSve"
   },
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pgGODxHsHDLE"
   },
   "outputs": [],
   "source": [
    "def get_score(sentence, model, tokenizer):\n",
    "    tokenize_input = tokenizer.tokenize(sentence)\n",
    "    tensor_input = torch.tensor([tokenizer.convert_tokens_to_ids(tokenize_input)])\n",
    "    predictions = model(tensor_input)\n",
    "    # print(predictions[0])\n",
    "    # sentence_embedding, word_embedding = predictions[0], predictions[1]\n",
    "    loss_fct = torch.nn.CrossEntropyLoss()\n",
    "    loss = loss_fct(predictions[0].squeeze(),tensor_input.squeeze()).data \n",
    "    return math.exp(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vaCzvXPhJbH2",
    "outputId": "2400b51e-0e18-43a2-bc80-1eaa999335b2"
   },
   "outputs": [],
   "source": [
    "get_score(\"Can you get some bread?\", bertMaskedLM, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fvYzCBbFK7X7",
    "outputId": "2aee96cb-6894-449d-c059-550ffad0e5ea"
   },
   "outputs": [],
   "source": [
    "get_score(\"This is a correct sentence.\", bertMaskedLM, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "82QXPQBDK-2u",
    "outputId": "cccb71c5-0897-40a3-8fd7-b13268efbdaa"
   },
   "outputs": [],
   "source": [
    "get_score(\"Fungus skatebutter\", bertMaskedLM, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FmbMW-eHKHTn",
    "outputId": "88467ec5-b990-4539-e4d6-7f712a50d925"
   },
   "outputs": [],
   "source": [
    "get_score(\"Xuusgs ddd\", bertMaskedLM, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O2DwGvS9LEKc"
   },
   "source": [
    "A lower perplexity score suggests a sentence with a higher probability of being correct."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6zpo1ZO-12KB"
   },
   "source": [
    "### Contextual Sentence Embeddings\n",
    "\n",
    "Standard BERT sentence embeddings don't perform as well as special architectures built for sentence embeddings, such as BERT siamese networks. \n",
    "\n",
    "We will use the [UKPLab package \"sentence transformers\"](https://github.com/UKPLab/sentence-transformers\n",
    "), which implement such methods.\n",
    "\n",
    "- Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks (EMNLP 2019)\n",
    "- Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation (EMNLP 2020)\n",
    "- Augmented SBERT: Data Augmentation Method for Improving Bi-Encoders for Pairwise Sentence Scoring Tasks (NAACL 2021)\n",
    "- The Curse of Dense Low-Dimensional Information Retrieval for Large Index Sizes (arXiv 2020)\n",
    "\n",
    "Documentation for [SBERT](https://www.sbert.net/examples/applications/computing-embeddings/README.html).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZTm6r2TRHw1C",
    "outputId": "8d128313-ce71-4704-f285-a7ad7b0294cd"
   },
   "outputs": [],
   "source": [
    "!pip install sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 66,
     "referenced_widgets": [
      "7eb35e89b40d4160be485628ead4aebb",
      "6d1fe792c2734dd9894723b493b230eb",
      "df2e14a318d649b4b932419f30c2b762",
      "7d4a3a09605b4fed8627177ebf5f37c6",
      "f7407246c9ed4cc098b953b45c03a553",
      "c6a792a3b276421aa49ebed2d325d567",
      "ec004cc52bc548dab857bf1632d4f8a9",
      "f183bf7ec0814c0898ac0a1a2afc06db"
     ]
    },
    "id": "K432TawSH0mV",
    "outputId": "ebda2adc-202e-4f20-820d-d623e2da8c5f"
   },
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "model = SentenceTransformer('distilbert-base-nli-stsb-mean-tokens')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qtnxsyhoH1nk"
   },
   "outputs": [],
   "source": [
    "#Our sentences we like to encode\n",
    "sentences = ['This framework generates embeddings for each input sentence',\n",
    "    'Sentences are passed as a list of string.', \n",
    "    'The quick brown fox jumps over the lazy dog.']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0jspJw3pH24w"
   },
   "outputs": [],
   "source": [
    "#Sentences are encoded by calling model.encode()\n",
    "embeddings = model.encode(sentences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CDkn9wcaZ1mD",
    "outputId": "fbe7ea75-c5bc-4cd3-d9ff-ed2824ae3863"
   },
   "outputs": [],
   "source": [
    "#Print the embeddings\n",
    "for sentence, embedding in zip(sentences, embeddings):\n",
    "    print(\"Sentence:\", sentence)\n",
    "    print(\"Embedding:\", embedding)\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y-ha_CxEIB45"
   },
   "source": [
    "We recommend visiting the documentation to view other uses of these sentence embeddings, such as [clustering](https://www.sbert.net/examples/applications/clustering/README.html), [Paraphrase Mining](https://www.sbert.net/examples/applications/paraphrase-mining/README.html), [Semantic Search](https://www.sbert.net/examples/applications/semantic-search/README.html), [Retreive and Rank](https://www.sbert.net/examples/applications/retrieve_rerank/README.html), and even multi-modal [Image Search](https://www.sbert.net/examples/applications/image-search/README.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UM-2pRuUa5Os"
   },
   "source": [
    "## BERTology\n",
    "\n",
    "\"BERTology\" refers to the growing number of papers disecting BERT and it's inner workings. \n",
    "\n",
    "- [A Primer in BERTology: What we know about how BERT works](https://arxiv.org/abs/2002.12327)\n",
    "\n",
    "- [Huggingface Transformers page on BERTology](https://huggingface.co/transformers/bertology.html)\n",
    "\n",
    "### Visualizing BERT Attention\n",
    "\n",
    "Just as we visualized our seq2seq attention, we can also visualize the tttention layers to understand BERT.\n",
    "\n",
    "**We highly recommend this visual blog on [explaining Transformers](http://jalammar.github.io/explaining-transformers/\n",
    ") and this visual blog on [visualising hidden states](http://jalammar.github.io/hidden-states/\n",
    ") before proceeding**\n",
    "\n",
    "The package \"bertviz\" [(GitHub link)](https://github.com/jessevig/bertviz) allows us to look into specific layers of multiple Transformers based models. \n",
    "\n",
    "This [blog post](https://towardsdatascience.com/deconstructing-bert-part-2-visualizing-the-inner-workings-of-attention-60a16d86b5c1) walks through the architecture of BERT and what exactly we can visualise. If you intend on using the BERT visualisations we recommend going through the blog post.\n",
    "\n",
    "We also point these two other visualisations tools for such models:\n",
    "- [Seq2Seq-Vis](http://seq2seq-vis.io/), [GitHub link](https://github.com/HendrikStrobelt/Seq2Seq-Vis)\n",
    "- [exBERT](https://exbert.net/), [GitHub link](https://github.com/bhoov/exbert)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "p-26bESCtirt",
    "outputId": "61f9d14b-7704-4c4f-eaf6-4ae803dd3621"
   },
   "outputs": [],
   "source": [
    "!pip install bertviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 262,
     "referenced_widgets": [
      "bf789bdff37a46faa45523e9eee49470",
      "b35dc95b8b874f39914c9d75101c87f7",
      "f3ced7fcad6d4bf785df92e435d7462c",
      "20471f637ccb43d0898b1d9d0b5fc87f",
      "955ebd5885fa483c8ce0d13d73fa0e69",
      "25a24033128540e4aded8e288e652209",
      "ed890a9a77b04c38ac89781ca89a55d7",
      "af0bd17be80c44a5b55fcd144a9067db",
      "0e777952989246c59ba83367274418d9",
      "1acf62efa6f0462e95ae7463650d5041",
      "fb527f6e7e384dc4816aa9bad433431e",
      "589ec9eac38e4ba485618558fe8ef475",
      "44a83d6440fd43468fa5abcb9f74573b",
      "6d33db5b406d4ebe9617d3b0385a8dd3",
      "8214f1afef1b4cb6a36f08b823375100",
      "77e51d8d225b47adb04e8fb5b0f5242b",
      "661f55db7f62465fbcffc4ebdb6005b2",
      "4eeaba199cbe4e0c9041079f50165b1d",
      "e50a63a1e3744f219b4f3d1fd1c7e229",
      "703b60bee7b34c8ea6fd3d8a416f272d",
      "2eadd7096d7d43ea9b6f2c608db0f3e1",
      "adedfff33f4c4154bd7714c86de319db",
      "738aecb0dbbd4ebe814aa00ee32f2208",
      "f09b87d956b744659aa87bd870b31156",
      "25a72e7b9e284386b33d996f3b5c76f6",
      "5cc2e9386c304304a97db2c9cb1d2408",
      "2fd31ea2dccc47a283d5115870d52b05",
      "2b779ae85e8f4830959ba2cea7d791ed",
      "d0d6290dd60d491b8e5a18b7a376bf8a",
      "32759be2bf6b4e0a8618c27885c3afb8",
      "a2ec074fb5ae4c9abead555edc752b3b",
      "c5db407c84914faa9992e9c90132f847",
      "9045da8b510044fb82e64e7071b19e9b",
      "bab4e42cefa4477ab18c836dab2e766a",
      "013eb02834c64a6c808db6a422a855e1",
      "ee35ba1d55614d959f9de73e8d68e3fc",
      "e80dec7f9165461b92d4f60d92e1c0d5",
      "9cbd92559d164ce69196c78d908d8b20",
      "07ed812a666c42fea2ffa35f127f51f1",
      "8c9002e562b44b16b41da22cf04ac61f"
     ]
    },
    "id": "48CAczKlt0Ej",
    "outputId": "38c5d5cf-8744-4b9b-fdfe-33d256664e76"
   },
   "outputs": [],
   "source": [
    "# Load model and retrieve attention weights\n",
    "\n",
    "from bertviz import head_view, model_view\n",
    "from transformers import BertTokenizer, BertModel\n",
    "\n",
    "model_version = 'bert-base-uncased'\n",
    "do_lower_case = True\n",
    "model = BertModel.from_pretrained(model_version, output_attentions=True)\n",
    "tokenizer = BertTokenizer.from_pretrained(model_version, do_lower_case=do_lower_case)\n",
    "sentence_a = \"The cat sat on the mat\"\n",
    "sentence_b = \"The cat lay on the rug\"\n",
    "inputs = tokenizer.encode_plus(sentence_a, sentence_b, return_tensors='pt', add_special_tokens=True)\n",
    "input_ids = inputs['input_ids']\n",
    "token_type_ids = inputs['token_type_ids']\n",
    "attention = model(input_ids, token_type_ids=token_type_ids)[-1]\n",
    "sentence_b_start = token_type_ids[0].tolist().index(1)\n",
    "input_id_list = input_ids[0].tolist() # Batch index 0\n",
    "tokens = tokenizer.convert_ids_to_tokens(input_id_list) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rffjP27Jt26g"
   },
   "source": [
    "## Model View\n",
    "The model view gives a birds-eye view of attention across all of the layers (rows) and heads (columns) in the model. In this case we are showing *bert-base*, which has 12 layers and 12 heads (zero-indexed). There are several GIFs on the [bertviz GitHub like this one](https://github.com/jessevig/bertviz):\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/UChicago-Computational-Content-Analysis/Homework-Notebooks/main/week-7/img/model-view-noscroll.gif\" alt=\"https://raw.githubusercontent.com/UChicago-Computational-Content-Analysis/Homework-Notebooks/main/week-7/img/model-view-noscroll.gif\" style=\"width:500px\">\n",
    "\n",
    "### Usage\n",
    "* **Click** on any **cell** for a detailed view of attention for the associated attention head.\n",
    "* Then **hover** over any **token** on the left side of detail view to filter the attention from that token.\n",
    "* The lines show the attention from each token (left) to every other token (right). Darker lines indicate higher attention weights.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "nT4T5bdEt54v",
    "outputId": "b7219861-ea3c-409e-8c7b-b2281ba85294"
   },
   "outputs": [],
   "source": [
    "model_view(attention, tokens, sentence_b_start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bYs0L8Ftt_Hu"
   },
   "source": [
    "## Head View\n",
    "The head view visualizes attention for one or more attention heads in the same layer. It is based on the excellent Tensor2Tensor visualization tool by Llion Jones. Here is the GIF from the [bertviz GitHub](https://github.com/jessevig/bertviz):\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/UChicago-Computational-Content-Analysis/Homework-Notebooks/main/week-7/img/head-view.gif\" alt=\"https://raw.githubusercontent.com/UChicago-Computational-Content-Analysis/Homework-Notebooks/main/week-7/img/head-view.gif\" style=\"width:500px\">\n",
    "\n",
    "### Usage\n",
    "* **Hover** over any **token** on the left/right side of the visualization to filter attention from/to that token. The colors correspond to different attention heads.\n",
    "* **Double-click** on any of the **colored tiles** at the top to filter to the corresponding attention head.\n",
    "* **Single-click** on any of the **colored tiles** to toggle selection of the corresponding attention head. \n",
    "* **Click** on the **Layer** drop-down to change the model layer (zero-indexed).\n",
    "* The lines show the attention from each token (left) to every other token (right). Darker lines indicate higher attention weights. When multiple heads are selected, the attention weights are overlaid on one another. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 408
    },
    "id": "6eqw4PQnuAcX",
    "outputId": "47eabf64-db25-4b77-a887-6fb04978660b"
   },
   "outputs": [],
   "source": [
    "head_view(attention, tokens, sentence_b_start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uyvp97I4uEj-"
   },
   "source": [
    "## Neuron View\n",
    "The attention-head view visualizes attention, as well as query and key values, in a particuler attention head.\n",
    "\n",
    "**NOTE: This visualization requires Chrome when run in Colab**. Here is the GIF:\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/UChicago-Computational-Content-Analysis/Homework-Notebooks/main/week-7/img/neuron-view-dark.gif\" alt=\"https://raw.githubusercontent.com/UChicago-Computational-Content-Analysis/Homework-Notebooks/main/week-7/img/neuron-view-dark.gif\" style=\"width:500px\">\n",
    "\n",
    "\n",
    "### Usage\n",
    "* **Hover** over any of the tokens on the left side of the visualization to filter attention from that token.\n",
    "* Then **click** on the **plus** icon that is revealed when hovering. This shows the query vectors, key vectors, and intermediate computations for the attention weights (<span style='color:blue'>blue</font>=positive, <span style='color:#ff6318'>orange</font>=negative).  \n",
    "* Once in the expanded view, **hover** over any other **token** on the left to see the associated attention computations.\n",
    "* **Click** on the **Layer** or **Head** drop-downs to change the model layer or head (zero-indexed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 546
    },
    "id": "MfNTtsiquD8K",
    "outputId": "6d43c8d1-6d53-4bda-c1af-bb528a5322aa"
   },
   "outputs": [],
   "source": [
    "from bertviz.transformers_neuron_view import BertModel, BertTokenizer\n",
    "from bertviz.neuron_view import show\n",
    "\n",
    "model = BertModel.from_pretrained(model_version, output_attentions=True)\n",
    "tokenizer = BertTokenizer.from_pretrained(model_version, do_lower_case=do_lower_case)\n",
    "model_type = 'bert'\n",
    "show(model, model_type, tokenizer, sentence_a, sentence_b, layer=2, head=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q35FQTJ6BEE8"
   },
   "source": [
    "These visualisations essentially allow us to see what attention values are assigned between a pair of sentences within the BERT model.\n",
    "\n",
    "There is also a package called [`ecco`](https://github.com/jalammar/ecco) with many different NLP interactive visualizations, including attention weights. Thanks to Austin Kozlowski for the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ecco"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load `distillgpt2`, another common base model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ecco\n",
    "lm = ecco.from_pretrained('distilgpt2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are two prediction tasks for GPT-2 and a visualization of what words in the prompt GPT-2 paid the most attention to in its prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"If you go to Washington DC, be sure to visit the\"\n",
    "# With ['ig', 'grad_x_input'] we are calculating Integrated Gradients and Input X Gradient saliency attributions\n",
    "output2 = lm.generate(text, generate=2, do_sample=False,beam_size=3, attribution=['grad_x_input'])\n",
    "output2.primary_attributions(attr_method='grad_x_input', style=\"detailed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"If you go to Paris, be sure to visit the\"\n",
    "# With ['ig', 'grad_x_input'] we are calculating Integrated Gradients and Input X Gradient saliency attributions\n",
    "output2 = lm.generate(text, generate=2, do_sample=False,beam_size=3, attribution=['grad_x_input'])\n",
    "output2.primary_attributions(attr_method='grad_x_input', style=\"detailed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the word \"White\" was predicted in large part from \"Washington\" and \"DC\", while \"French\" was predicted in large part from \"Paris\"! Nifty."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <font color=\"red\">*Exercise 2*</font>\n",
    "\n",
    "<font color=\"red\">Construct cells immediately below this that visualize your BERT model with the `bertviz` and/or `ecco` packages. There is also a similar tool called [`exbert`](https://huggingface.co/exbert/?model=distilgpt2&modelKind=bidirectional&sentence=The%20girl%20ran%20to%20a%20local%20pub%20to%20escape%20the%20din%20of%20her%20city.&layer=0&heads=..0,1,2,3,4,5,6,7,8,9,10,11&threshold=0.7&tokenInd=null&tokenSide=null&maskInds=..&hideClsSep=true) that may be of interest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pOrLpn_G-ZMp"
   },
   "source": [
    "### Language Modeling with 'pipelines'\n",
    "We will now use the popular Transformers package default 'pipelines' feature to do language modeling.\n",
    "Language modeling is the task of fitting a model to a corpus, which can be domain specific. All popular\n",
    "transformer-based models are trained using a variant of language modeling, e.g., BERT with masked language modeling,\n",
    "GPT-2/3 with causal language modeling.\n",
    "\n",
    "Language modeling can be useful beyond pretraining as well, for example to shift the model distribution to be\n",
    "domain-specific: using a language model trained over a very large corpus, and then fine-tuning it to a news dataset or\n",
    "on scientific papers e.g. [LysandreJik/arxiv-nlp](https://huggingface.co/lysandre/arxiv-nlp)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-zwgOcdZ-ZMp"
   },
   "source": [
    "#### Masked Language Modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4r5bR8hM-ZMp"
   },
   "source": [
    "Masked language modeling is the task of masking tokens in a sequence with a masking token, and prompting the model to\n",
    "fill that mask with an appropriate token. This allows the model to attend to both the right context (tokens on the\n",
    "right of the mask) and the left context (tokens on the left of the mask). Such a training creates a strong basis for\n",
    "downstream tasks requiring bi-directional context, such as SQuAD (question answering, see [Lewis, Lui, Goyal et al.](https://arxiv.org/abs/1910.13461), part 4.2).\n",
    "\n",
    "Here is an example of using pipelines to replace a mask from a sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GBTkRCoE-ZMq"
   },
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "nlp = pipeline(\"fill-mask\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jDiURWsT-ZMq"
   },
   "source": [
    "This outputs the sequences with the mask filled, the confidence score, and the token id in the tokenizer vocabulary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ayEq2LE2-ZMq",
    "outputId": "4f3ddaf8-e050-4710-f204-d8c27931720b"
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "pprint(nlp(f\"HuggingFace is creating a {nlp.tokenizer.mask_token} that the community uses to solve NLP tasks.\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hUN9TACu-ZMq"
   },
   "source": [
    "Here is an example of doing masked language modeling using a model and a tokenizer. The process as follows:\n",
    "\n",
    "1. Instantiate a tokenizer and a model from the checkpoint name. The model is identified as a DistilBERT model and\n",
    "   loads it with the weights stored in the checkpoint.\n",
    "2. Define a sequence with a masked token, placing the `tokenizer.mask_token` instead of a word.\n",
    "3. Encode that sequence into a list of IDs and find the position of the masked token in that list.\n",
    "4. Retrieve the predictions at the index of the mask token: this tensor has the same size as the vocabulary, and the\n",
    "   values are the scores attributed to each token. The model gives higher scores to tokens it deems probable in that\n",
    "   context.\n",
    "5. Retrieve the top 5 tokens using the PyTorch `topk` or TensorFlow `top_k` methods.\n",
    "6. Replace the mask token by the tokens and print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You may need transformers installed from source to get the distilled models.\n",
    "# ! pip install -U git+git://github.com/huggingface/transformers.git -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ihYScqSt-ZMq"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelWithLMHead, AutoTokenizer\n",
    "import torch\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-cased\")\n",
    "model = AutoModelWithLMHead.from_pretrained(\"distilbert-base-cased\")\n",
    "sequence = f\"Distilled models are smaller than the models they mimic. Using them instead of the large versions would help {tokenizer.mask_token} our carbon footprint.\"\n",
    "input = tokenizer.encode(sequence, return_tensors=\"pt\")\n",
    "mask_token_index = torch.where(input == tokenizer.mask_token_id)[1]\n",
    "token_logits = model(input).logits\n",
    "mask_token_logits = token_logits[0, mask_token_index, :]\n",
    "top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hVBAG5Al-ZMr"
   },
   "source": [
    "This prints five sequences, with the top 5 tokens predicted by the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NyBSdXGp-ZMr",
    "outputId": "c2813939-393b-43d7-eee5-27681e8b617d"
   },
   "outputs": [],
   "source": [
    "for token in top_5_tokens:\n",
    "    print(sequence.replace(tokenizer.mask_token, tokenizer.decode([token])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xxGiKa1h-ZMr"
   },
   "source": [
    "#### Causal Language Modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qt7FighC-ZMr"
   },
   "source": [
    "Causal language modeling is the task of predicting the token following a sequence of tokens. In this situation, the\n",
    "model only attends to the left context (tokens on the left of the mask). Such a training is particularly interesting\n",
    "for generation tasks.\n",
    "\n",
    "Usually, the next token is predicted by sampling from the logits of the last hidden state the model produces from the\n",
    "input sequence.\n",
    "\n",
    "Here is an example of using the tokenizer and model and leveraging the\n",
    "`PreTrainedModel.top_k_top_p_filtering` method to sample the next token following an input sequence\n",
    "of tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h8mfXD4L-ZMs"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelWithLMHead, AutoTokenizer, top_k_top_p_filtering\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "model = AutoModelWithLMHead.from_pretrained(\"gpt2\")\n",
    "sequence = f\"Hugging Face is based in DUMBO, New York City, and \"\n",
    "input_ids = tokenizer.encode(sequence, return_tensors=\"pt\")\n",
    "# get logits of last hidden state\n",
    "next_token_logits = model(input_ids).logits[:, -1, :]\n",
    "# filter\n",
    "filtered_next_token_logits = top_k_top_p_filtering(next_token_logits, top_k=50, top_p=1.0)\n",
    "# sample\n",
    "probs = F.softmax(filtered_next_token_logits, dim=-1)\n",
    "next_token = torch.multinomial(probs, num_samples=1)\n",
    "generated = torch.cat([input_ids, next_token], dim=-1)\n",
    "resulting_string = tokenizer.decode(generated.tolist()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kpULOLx0-ZMs"
   },
   "source": [
    "This outputs a (hopefully) coherent next token following the original sequence, which in our case is the word **has**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fHzm0TFJ-ZMs",
    "outputId": "328d3181-2c0c-466b-ca06-2ae15ebc382f"
   },
   "outputs": [],
   "source": [
    "print(resulting_string)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <font color=\"red\">*Exercise 3*</font>\n",
    "\n",
    "<font color=\"red\">In the cells immediately following, use the pipeline functions or the word or sentence vector functions (e.g., similarity) to explore the social game underlying the production and meaning of texts associated with your final project. This may be similar to what you did in Week 5, but try to make use of new things you've learned since then. How does BERT help you gain insight regarding your research question that is similar and different from prior methods?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}