contextualreps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bringing contextual word representations into your models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"Christopher Potts\"\n",
    "__version__ = \"CS224u, Stanford, Spring 2020\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "1. [Overview](#Overview)\n",
    "1. [General set-up](#General-set-up)\n",
    "1. [Hugging Face BERT interfaces](#Hugging-Face-BERT-interfaces)\n",
    "  1. [Hugging Face BERT set-up](#Hugging-Face-BERT-set-up)\n",
    "  1. [Hugging Face BERT basics](#Hugging-Face-BERT-basics)\n",
    "  1. [BERT featurization with Hugging Face](#BERT-featurization-with-Hugging-Face)\n",
    "    1. [Simple feed-forward experiment](#Simple-feed-forward-experiment)\n",
    "    1. [A feed-forward experiment with the sst module](#A-feed-forward-experiment-with-the-sst-module)\n",
    "    1. [An RNN experiment with the sst module](#An-RNN-experiment-with-the-sst-module)\n",
    "  1. [BERT fine-tuning with Hugging Face](#BERT-fine-tuning-with-Hugging-Face)\n",
    "1. [Using ELMo](#Using-ELMo)\n",
    "  1. [ELMO Allen NLP set-up](#ELMO-Allen-NLP-set-up)\n",
    "  1. [ELMo featurization](#ELMo-featurization)\n",
    "    1. [ELMo featurization for an RNN](#ELMo-featurization-for-an-RNN)\n",
    "    1. [Using the SST experiment framework with ELMo](#Using-the-SST-experiment-framework-with-ELMo)\n",
    "  1. [ELMo fine-tuning](#ELMo-fine-tuning)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This notebook provides a basic introduction to using pre-trained [BERT](https://github.com/google-research/bert) and [ELMo](https://allennlp.org/elmo) representations. It is meant as a practical companion to our lecture on contextual word representations. The goal of this notebook is just to help you use these representations in your own work. The BERT and ELMo teams have done amazing work to make these resources available to the community. Many projects can benefit from them, so it is probably worth your time to experiment.\n",
    "\n",
    "This notebook should be considered an experimental extension to the regular course materials. It has some special requirements – libraries and data files – that are not part of the core requirements for this repository. All these tools are very new and being updated frequently, so you might need to do some fiddling to get all of this to work. As I said, though, it's probably worth the effort!\n",
    "\n",
    "A number of the experiments in this notebook are resource intensive. I've included timing information for the expensive steps, to give you a sense for how long things are likely to take. I ran this notebook on a 2015 iMac with a 4 GHz Intel Core i7 CPU (no GPU involved) and 32GB of memory, and none of the steps seemed to strain the system. If you run this notebook on a GPU that PyTorch can work with, everything will be much faster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General set-up\n",
    "\n",
    "The following are requirements that you'll already have met if you've been working in this repository. As you can see, we'll use the [Stanford Sentiment Treebank](sst_01_overview.ipynb) for illustrations, and we'll try out a few different deep learning models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sst\n",
    "from torch_shallow_neural_classifier import TorchShallowNeuralClassifier\n",
    "from torch_rnn_classifier import TorchRNNClassifier, TorchRNNClassifierModel\n",
    "from torch_rnn_classifier import TorchRNNClassifier\n",
    "from sklearn.metrics import classification_report\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set all the random seeds for reproducibility. Only the\n",
    "# system and torch seeds are relevant for this notebook.\n",
    "\n",
    "utils.fix_random_seeds()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "SST_HOME = os.path.join(\"data\", \"trees\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hugging Face BERT interfaces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hugging Face BERT set-up\n",
    "\n",
    "To install this library, run\n",
    "\n",
    "```pip install transformers```\n",
    "\n",
    "I've tested this code with versions 2.4 and 2.5 of `transformers`. Try to get 2.5. It requires `pip >= 20` and, I think, a version of [Rust](https://www.rust-lang.org) at least as high as 1.21.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0309 20:28:13.657829 140735854211968 file_utils.py:41] PyTorch version 1.3.0 available.\n",
      "I0309 20:28:13.658627 140735854211968 file_utils.py:57] TensorFlow version 2.0.0 available.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from transformers import BertModel, BertTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `transformers` library does a lot of logging. To avoid ending up with a cluttered notebook, I am changing the logging level. You might want to skip this as you scale up to building production systems, since the logging is very good – it gives you a lot of insights into what the models and code are doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logger = logging.getLogger()\n",
    "logger.level = logging.ERROR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hugging Face BERT basics\n",
    "\n",
    "To start, let's get a feel for the basic API that `transformers` provides. The first step is specifying the pretrained parameters we'll be using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_weights_name = 'bert-base-cased'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are lots other options for pretrained weights. See [this section of the project README.md](https://github.com/huggingface/transformers#quick-tour) for a good overview and code that documents how these weights align with different Transformer model classes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we specify a tokenizer and a model that match both each other and our choice of pretrained weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_tokenizer = BertTokenizer.from_pretrained(hf_weights_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_model = BertModel.from_pretrained(hf_weights_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's illuminating to see what the tokenizer does to example texts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_example_texts = [\n",
    "    \"Encode sentence 1. [SEP] And sentence 2!\",\n",
    "    \"Bert knows Snuffleupagus\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `encode` method maps individual strings to indices into the underlying embedding used by the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[101, 13832, 13775, 5650, 122, 119, 102, 1262, 5650, 123, 106, 102]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ex0_ids = hf_tokenizer.encode(hf_example_texts[0], add_special_tokens=True)\n",
    "\n",
    "ex0_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can get a better feel for what these representations are like by mapping the indices back to \"words\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['[CLS]',\n",
       " 'En',\n",
       " '##code',\n",
       " 'sentence',\n",
       " '1',\n",
       " '.',\n",
       " '[SEP]',\n",
       " 'And',\n",
       " 'sentence',\n",
       " '2',\n",
       " '!',\n",
       " '[SEP]']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_tokenizer.convert_ids_to_tokens(ex0_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For modeling, we will often need to pad (and perhaps truncate) token lists so that we can work with fixed-dimensional tensors: The `batch_encode_plus` has a lot of options for doing this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_example_ids = hf_tokenizer.batch_encode_plus(\n",
    "    hf_example_texts, \n",
    "    add_special_tokens=True,\n",
    "    return_attention_mask=True,\n",
    "    pad_to_max_length=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_example_ids.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `token_type_ids` is used for multi-text inputs like NLI. The `'input_ids'` field gives the indices for each of the two examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[101, 13832, 13775, 5650, 122, 119, 102, 1262, 5650, 123, 106, 102],\n",
       " [101, 15035, 3520, 156, 14787, 13327, 4455, 28026, 1116, 102, 0, 0]]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_example_ids['input_ids']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For fine-tuning, we want to avoid attending to padded tokens. The `'attention_mask'` captures the needed mask, which we'll be able to feed directly to the pretrained BERT model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_example_ids['attention_mask']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can run these indices and masks through the pretrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_hf_example = torch.tensor(hf_example_ids['input_ids'])\n",
    "X_hf_example_mask = torch.tensor(hf_example_ids['attention_mask'])\n",
    "\n",
    "with torch.no_grad():    \n",
    "    hf_final_hidden_states, cls_output = hf_model(\n",
    "        X_hf_example, attention_mask=X_hf_example_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "BERT representations are pretty large – this shows 2 examples, with the second padded to the length of the larger one in the batch (12). The individual representations have dimensionality 768."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 12, 768])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_final_hidden_states.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Those are all the essential ingredients for working with these parameters in Hugging Face. Of course, the library has a lot of other functionality, but the above suffices to featurize and to fine tune."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BERT featurization with Hugging Face\n",
    "\n",
    "To start, we'll use the Hugging Face interfaces just to featurize examples to create inputs to a separate model. In this setting, the BERT parameters are frozen. The heart of this approach is the following featurizer, which flattens an SST tree into a string, tokenizes it, and calculates its hidden representations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hugging_face_bert_phi(tree):\n",
    "    s = \" \".join(tree.leaves())\n",
    "    input_ids = hf_tokenizer.encode(s, add_special_tokens=True)\n",
    "    X = torch.tensor([input_ids])\n",
    "    with torch.no_grad():\n",
    "        final_hidden_states, cls_output = hf_model(X)\n",
    "        return final_hidden_states.squeeze(0).numpy() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Simple feed-forward experiment\n",
    "\n",
    "For a simple feed-forward experiment, we can get the representation of the `[CLS]` tokens and use them as the inputs to a shallow neural network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hugging_face_bert_classifier_phi(tree):\n",
    "    reps = hugging_face_bert_phi(tree)\n",
    "    #return reps.mean(axis=0)  # Another good, easy option.\n",
    "    return reps[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is very much like what we [summed the GloVe representations of these examples](sst_03_neural_networks.ipynb#Distributed-representations-as-features), but now the individual word representations are different depending on the context in which they appear."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we read in the SST train and dev portions as a lists of `(tree, label)` pairs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_train = list(sst.train_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "hf_dev = list(sst.dev_reader(SST_HOME, class_func=sst.ternary_class_func))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split the input/output pairs out into separate lists:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_hf_tree_train, y_hf_train = zip(*hf_train)\n",
    "\n",
    "X_hf_tree_dev, y_hf_dev = zip(*hf_dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next step, we featurize all of the examples. These steps are likely to be the slowest in these experiments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 22min 28s, sys: 5.87 s, total: 22min 34s\n",
      "Wall time: 5min 38s\n"
     ]
    }
   ],
   "source": [
    "%time X_hf_train = [hugging_face_bert_classifier_phi(tree) for tree in X_hf_tree_train]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2min 51s, sys: 681 ms, total: 2min 52s\n",
      "Wall time: 43.1 s\n"
     ]
    }
   ],
   "source": [
    "%time X_hf_dev = [hugging_face_bert_classifier_phi(tree) for tree in X_hf_tree_dev]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that all the examples are featurized, we can fit a model and evaluate it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_mod = TorchShallowNeuralClassifier(max_iter=100, hidden_dim=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 100 of 100; error is 3.4869229197502136"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 45s, sys: 1.44 s, total: 1min 47s\n",
      "Wall time: 15.4 s\n"
     ]
    }
   ],
   "source": [
    "%time _ = hf_mod.fit(X_hf_train, y_hf_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_preds = hf_mod.predict(X_hf_dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.710     0.640     0.673       428\n",
      "     neutral      0.331     0.218     0.263       229\n",
      "    positive      0.676     0.858     0.756       444\n",
      "\n",
      "    accuracy                          0.640      1101\n",
      "   macro avg      0.572     0.572     0.564      1101\n",
      "weighted avg      0.617     0.640     0.621      1101\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(y_hf_dev, hf_preds, digits=3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A feed-forward experiment with the sst module\n",
    "\n",
    "It is straightforward to conduct experiments like the above using `sst.experiment`, which will enable you to do a wider range of experiments without writing or copy-pasting a lot of code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_hf_shallow_network(X, y):\n",
    "    mod = TorchShallowNeuralClassifier(\n",
    "        max_iter=100, hidden_dim=300)\n",
    "    mod.fit(X, y)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 100 of 100; error is 2.8084703385829926"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.670     0.764     0.714       428\n",
      "     neutral      0.333     0.205     0.254       229\n",
      "    positive      0.714     0.759     0.736       444\n",
      "\n",
      "    accuracy                          0.646      1101\n",
      "   macro avg      0.572     0.576     0.568      1101\n",
      "weighted avg      0.618     0.646     0.627      1101\n",
      "\n",
      "CPU times: user 27min 11s, sys: 7.76 s, total: 27min 19s\n",
      "Wall time: 6min 39s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "_ = sst.experiment(\n",
    "    SST_HOME,\n",
    "    hugging_face_bert_classifier_phi,\n",
    "    fit_hf_shallow_network,\n",
    "    train_reader=sst.train_reader, \n",
    "    assess_reader=sst.dev_reader, \n",
    "    class_func=sst.ternary_class_func,\n",
    "    vectorize=False)  # Pass in the BERT hidden state directly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### An RNN experiment with the sst module\n",
    "\n",
    "We can also use BERT representations as the input to an RNN. There is just one key change from how we used these models before:\n",
    "\n",
    "* Previously, we would feed in lists of tokens, and they would be converted to indices into a fixed embedding space. This presumes that all words have the same representation no matter what their context is. \n",
    "\n",
    "* With BERT, we skip the embedding entirely and just feed in lists of BERT vectors, which means that the same word can be represented in different ways.\n",
    "\n",
    "`TorchRNNClassifier` supports this via `use_embedding=False`. In turn, you needn't supply a vocabulary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_hf_rnn(X, y):\n",
    "    mod = TorchRNNClassifier(\n",
    "        vocab=[],\n",
    "        max_iter=50, \n",
    "        hidden_dim=50,\n",
    "        use_embedding=False)  # Pass in the BERT hidden states directly!\n",
    "    mod.fit(X, y)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 50 of 50; error is 0.19242170546203852"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.709     0.706     0.707       428\n",
      "     neutral      0.322     0.249     0.281       229\n",
      "    positive      0.715     0.802     0.756       444\n",
      "\n",
      "    accuracy                          0.649      1101\n",
      "   macro avg      0.582     0.585     0.581      1101\n",
      "weighted avg      0.631     0.649     0.638      1101\n",
      "\n",
      "CPU times: user 54min 10s, sys: 2min 53s, total: 57min 4s\n",
      "Wall time: 14min 26s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "_ = sst.experiment(\n",
    "    SST_HOME,\n",
    "    hugging_face_bert_phi,\n",
    "    fit_hf_rnn,\n",
    "    train_reader=sst.train_reader, \n",
    "    assess_reader=sst.dev_reader, \n",
    "    class_func=sst.ternary_class_func,\n",
    "    vectorize=False)  # Pass in the BERT hidden states directly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BERT fine-tuning with Hugging Face\n",
    "\n",
    "The above experiments are quite successful – BERT gives us a reliable boost compared to other methods we've explored for the SST task. However, we might expect to do even better if we fine-tune the BERT parameters as part of fitting our SST classifier. To do that, we need to incorporate the Hugging Face BERT model into our classifier. This too is quite straightforward.\n",
    "\n",
    "The most important step is to create an `nn.Module` subclass that has, for its parameters, both the BERT model and parameters for our own classifier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HfBertClassifierModel(nn.Module):\n",
    "    def __init__(self, n_classes, weights_name='bert-base-cased'):\n",
    "        super().__init__()\n",
    "        self.n_classes = n_classes\n",
    "        self.weights_name = weights_name\n",
    "        self.bert = BertModel.from_pretrained(self.weights_name)\n",
    "        self.hidden_dim = self.bert.embeddings.word_embeddings.embedding_dim\n",
    "        # The only new parameters -- the classifier layer:\n",
    "        self.W = nn.Linear(self.hidden_dim, self.n_classes)\n",
    "        \n",
    "    def forward(self, X):\n",
    "        \"\"\"Here, `X` is an np.array in which each element is a pair \n",
    "        consisting of an index into the BERT embedding and a 1 or 0\n",
    "        indicating whether the token is masked. The `fit` method will \n",
    "        train all these parameters against a softmax objective.\n",
    "        \n",
    "        \"\"\"\n",
    "        indices = X[: , 0, : ]\n",
    "        # Type conversion, since the base class insists on\n",
    "        # casting this as a FloatTensor, but we ned Long\n",
    "        # for `bert`.\n",
    "        indices = indices.long()\n",
    "        mask = X[: , 1, : ]      \n",
    "        (final_hidden_states, cls_output) = self.bert(\n",
    "            indices, attention_mask=mask)       \n",
    "        return self.W(cls_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the training and prediction interface, we can somewhat opportunistically subclass `TorchShallowNeuralClassifier` so that we don't have to write any of our own data-handling, training, or prediction code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HfBertClassifier(TorchShallowNeuralClassifier):\n",
    "    def __init__(self, weights_name, *args, **kwargs):\n",
    "        self.weights_name = weights_name\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(self.weights_name)\n",
    "        super().__init__(*args, **kwargs)\n",
    "        \n",
    "    def define_graph(self):\n",
    "        \"\"\"This method is used by `fit`. We override it here to use our\n",
    "        new BERT-based graph.\n",
    "        \n",
    "        \"\"\"\n",
    "        bert = HfBertClassifierModel(\n",
    "            self.n_classes_, weights_name=self.weights_name)\n",
    "        bert.train()\n",
    "        return bert\n",
    "    \n",
    "    def encode(self, X, max_length=None):\n",
    "        \"\"\"The `X` is a list of strings. We use the model's tokenizer\n",
    "        to get the indices and mask information.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        list of [index, mask] pairs, where index is an int and mask\n",
    "        is 0 or 1.\n",
    "        \n",
    "        \"\"\"\n",
    "        data = self.tokenizer.batch_encode_plus(\n",
    "            X, \n",
    "            max_length=max_length,\n",
    "            add_special_tokens=True, \n",
    "            pad_to_max_length=True,\n",
    "            return_attention_mask=True)\n",
    "        indices = data['input_ids']\n",
    "        mask = data['attention_mask']\n",
    "        return [[i, m] for i, m in zip(indices, mask)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a self-contained illustration, starting from the raw SST data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_train = list(sst.train_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "hf_dev = list(sst.dev_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "X_hf_tree_train, y_hf_train = zip(*hf_train)\n",
    "\n",
    "X_hf_tree_dev, y_hf_dev = zip(*hf_dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our model has some standard fine-tuning parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_fine_tune_mod = HfBertClassifier(\n",
    "    'bert-base-cased', \n",
    "    batch_size=16, # Crucial; large batches will eat up all your memory!\n",
    "    max_iter=4, \n",
    "    eta=0.00002)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the `encode` method in `HfBertClassifier`. For generality, this is expecting string inputs, rather than trees, so we first flatten the trees:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_hf_str_train = [\" \".join(tree.leaves()) for tree in X_hf_tree_train]\n",
    "\n",
    "X_hf_str_dev = [\" \".join(tree.leaves()) for tree in X_hf_tree_dev]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can encode them; this step packs together the indices and mask information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_hf_indices_train = hf_fine_tune_mod.encode(X_hf_str_train)\n",
    "\n",
    "X_hf_indices_dev = hf_fine_tune_mod.encode(X_hf_str_dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training this model is resource intensive. Be patient – it will be worth the wait! (This experiment takes about 10 minutes on a machine with an NVIDIA RTX 2080 Max-Q GPU.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 4 of 4; error is 100.52562420163304"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 10h 16min 28s, sys: 38min 35s, total: 10h 55min 4s\n",
      "Wall time: 2h 49min 37s\n"
     ]
    }
   ],
   "source": [
    "%time _ = hf_fine_tune_mod.fit(X_hf_indices_train, y_hf_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, some predictions on the dev set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_fine_tune_preds = hf_fine_tune_mod.predict(X_hf_indices_dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.710     0.840     0.770       362\n",
      "     neutral      0.402     0.372     0.387       247\n",
      "    positive      0.854     0.770     0.810       492\n",
      "\n",
      "    accuracy                          0.704      1101\n",
      "   macro avg      0.655     0.661     0.655      1101\n",
      "weighted avg      0.705     0.704     0.702      1101\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(hf_fine_tune_preds, y_hf_dev, digits=3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is just one of the many possible ways to fine-tune BERT using our course modules or new modules you write. The crux of it is creating an `nn.Module` that combines the BERT parameters with your model's new parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using ELMo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ELMO Allen NLP set-up\n",
    "\n",
    "There are a number of ways to use pre-trained ELMo models. We'll use the simplest of the AllenNLP interfaces. Run the following to install [AllenNLP](https://allennlp.org):\n",
    "\n",
    "```sh\n",
    "pip install allennlp\n",
    "```\n",
    "I've tested this notebook with versions 0.8.0 and 0.9.0.\n",
    "\n",
    "Mac users: If your installation fails, make sure your Xcode tools are up to date by running `xcode-select --install`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex.\n"
     ]
    }
   ],
   "source": [
    "from allennlp.modules.elmo import Elmo, batch_to_ids\n",
    "from allennlp.commands.elmo import ElmoEmbedder\n",
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the following models, which will download from S3 to a local temp directory the first time you use them with `ElmoEmbedder` or `Elmo` as described below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_file_path = \"https://allennlp.s3.amazonaws.com/models/elmo/2x4096_512_2048cnn_2xhighway/\"\n",
    "\n",
    "options_file = elmo_file_path + \"elmo_2x4096_512_2048cnn_2xhighway_options.json\"\n",
    "\n",
    "weights_file = elmo_file_path + \"elmo_2x4096_512_2048cnn_2xhighway_weights.hdf5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more models:\n",
    "\n",
    "https://allennlp.org/elmo\n",
    "\n",
    "For additional details:\n",
    "\n",
    "https://github.com/allenai/allennlp/blob/master/allennlp/commands/elmo.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ELMo featurization\n",
    "\n",
    "As we did with BERT, we'll first use ELMo just to featurize examples for a separate model. The `ElmoEmbedder` facilitates this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_embedder = ElmoEmbedder(options_file, weights_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a copy of the SST in a useful format for this modeling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_train = list(sst.train_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "elmo_dev = list(sst.dev_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "X_elmo_tree_train, y_elmo_train = zip(*elmo_train)\n",
    "\n",
    "X_elmo_tree_dev, y_elmo_dev = zip(*elmo_dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ELMo interface requires tokenized input. I believe the tokenization scheme is the same as for the SST, so we can just use the leaves of those trees:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_elmo_toks_train = [tree.leaves() for tree in X_elmo_tree_train]\n",
    "\n",
    "X_elmo_toks_dev = [tree.leaves() for tree in X_elmo_tree_dev]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ELMo featurization for an RNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we create the representations for the train and dev sets; these steps are somewhat slow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 22min 14s, sys: 1min 55s, total: 24min 9s\n",
      "Wall time: 8min 57s\n"
     ]
    }
   ],
   "source": [
    "%time X_elmo_train_layers = list(elmo_embedder.embed_sentences(X_elmo_toks_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2min 51s, sys: 15.1 s, total: 3min 6s\n",
      "Wall time: 1min 9s\n"
     ]
    }
   ],
   "source": [
    "%time X_elmo_dev_layers = list(elmo_embedder.embed_sentences(X_elmo_toks_dev))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each member of `X_elmo_train_layers` has three dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3, 13, 1024)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_elmo_dev_layers[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each word (second dimension), there are three layers of length 1024. So ELMo representations are even larger than BERT ones!\n",
    "\n",
    "There are many ways we could combine the layers available for each word. Here, I'll use just the top layer; see [section 3.2 of the ELMo paper](https://www.aclweb.org/anthology/N18-1202/) for additional ideas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elmo_layer_reduce_top(elmo_vecs):\n",
    "    return [ex[-1] for ex in elmo_vecs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_elmo_train = elmo_layer_reduce_top(X_elmo_train_layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can fit an RNN as usual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_rnn = TorchRNNClassifier(\n",
    "    vocab=[],\n",
    "    max_iter=50,\n",
    "    use_embedding=False) # Pass in the ELMo hidden states directly!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 50 of 50; error is 0.7801527082920074"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 21min 4s, sys: 2min 14s, total: 23min 19s\n",
      "Wall time: 5min 59s\n"
     ]
    }
   ],
   "source": [
    "%time _ = elmo_rnn.fit(X_elmo_train, y_elmo_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation proceeds in the usual way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_elmo_dev = elmo_layer_reduce_top(X_elmo_dev_layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_rnn_preds = elmo_rnn.predict(X_elmo_dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.673     0.710     0.691       428\n",
      "     neutral      0.323     0.231     0.270       229\n",
      "    positive      0.711     0.777     0.743       444\n",
      "\n",
      "    accuracy                          0.638      1101\n",
      "   macro avg      0.569     0.573     0.568      1101\n",
      "weighted avg      0.616     0.638     0.624      1101\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(y_elmo_dev, elmo_rnn_preds, digits=3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using the SST experiment framework with ELMo\n",
    "\n",
    "To round things out, here's an example of how to use `sst.experiment` with ELMo, for more compact and maintainable experiment code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elmo_sentence_phi(tree):\n",
    "    vecs = elmo_embedder.embed_sentence(tree.leaves())\n",
    "    return vecs[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_elmo_rnn(X, y):\n",
    "    mod = TorchRNNClassifier(\n",
    "        vocab=[],\n",
    "        max_iter=50,\n",
    "        use_embedding=False)\n",
    "    mod.fit(X, y)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This step re-encodes all of the examples, so it will take a while before the model starts training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 50 of 50; error is 0.034750474384054545"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.690     0.687     0.689       428\n",
      "     neutral      0.264     0.183     0.216       229\n",
      "    positive      0.707     0.822     0.760       444\n",
      "\n",
      "    accuracy                          0.637      1101\n",
      "   macro avg      0.554     0.564     0.555      1101\n",
      "weighted avg      0.608     0.637     0.619      1101\n",
      "\n",
      "CPU times: user 3h 33min 19s, sys: 2min 49s, total: 3h 36min 9s\n",
      "Wall time: 54min 9s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "_ = sst.experiment(\n",
    "    SST_HOME,\n",
    "    elmo_sentence_phi,\n",
    "    fit_elmo_rnn,\n",
    "    train_reader=sst.train_reader, \n",
    "    assess_reader=sst.dev_reader, \n",
    "    class_func=sst.ternary_class_func,\n",
    "    vectorize=False)  # Pass in the ELMo hidden states directly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ELMo fine-tuning\n",
    "\n",
    "Fine-tuning ELMo proceeds in essentially the same way it did for BERT: we create an `nn.Module` that combines the parameters from ELMo with our task-specific parameters and then optimize everything on the new task. To illustrate, I'll define an RNN on top of the ELMo model using new subclasses of `TorchRNNClassifier` and `TorchRNNClassifierModel`.\n",
    "\n",
    "To start, let's get a feel for the primary interface, and then we'll write the classes that will allow us to use these components systematically.\n",
    "\n",
    "The interface to the ELMo parameters in this context is the class `Elmo`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo = Elmo(options_file, weights_file, num_output_representations=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model expects tokenized inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_example_texts = [\n",
    "    [\"Encode\", \"sentence\", \"1\", \".\"],\n",
    "    [\"ELMo\", \"knows\" \"Snuffleupagus\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ELMo model processes its tokens at the character-level, creating convolutional representations for the words from various character n-gram combinations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([259,  70, 111, 100, 112, 101, 102, 260, 261, 261, 261, 261, 261, 261,\n",
       "        261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261,\n",
       "        261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261,\n",
       "        261, 261, 261, 261, 261, 261, 261, 261])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elmo_character_ids = batch_to_ids(elmo_example_texts)\n",
    "\n",
    "# First word of the first example:\n",
    "elmo_character_ids[0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`elmo` embeds these at the word-level:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_embeddings = elmo(elmo_character_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`elmo_embeddings` is a dict. The value of `'elmo_representations'` is a list of tensors corresponding to each layer of the model. In other words, each tensor in the list is a complete representation of the example. The final element of the list is the final representation layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[[-0.7944,  0.0000,  0.0318,  ..., -0.0000, -0.8271, -0.0000],\n",
       "          [-0.0379,  0.0000,  0.0000,  ...,  0.0281, -0.0000, -0.2358],\n",
       "          [ 0.1434,  0.5358,  0.7767,  ..., -0.6500, -0.0777, -0.0000],\n",
       "          [-0.0000, -0.4965, -0.0000,  ..., -0.0000,  0.0722,  0.0000]],\n",
       " \n",
       "         [[ 0.0000, -0.0000,  0.0000,  ..., -0.8601, -0.0000, -0.0000],\n",
       "          [ 0.0000, -0.0000, -0.0000,  ..., -0.2244, -0.0000,  0.7476],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],\n",
       "        grad_fn=<MulBackward0>),\n",
       " tensor([[[-0.7944,  0.0000,  0.0000,  ..., -0.0000, -0.8271, -1.9283],\n",
       "          [-0.0000,  0.0000,  0.6947,  ...,  0.0000, -0.0000, -0.2358],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ..., -0.0000, -0.0000, -0.0000],\n",
       "          [-1.1810, -0.0000, -1.7275,  ..., -0.1605,  0.0722,  0.0000]],\n",
       " \n",
       "         [[ 0.4553, -0.0000,  0.0000,  ..., -0.8601, -0.0000, -0.5380],\n",
       "          [ 0.0000, -0.0000, -0.0832,  ..., -0.2244, -0.2797,  0.7476],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],\n",
       "        grad_fn=<MulBackward0>)]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elmo_embeddings['elmo_representations']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the representations we will be fine-tuning. There are many ways to combine use. In my simple illustration, I just take the top layer, as we did in the simpler featurization example above, but now keeping each word representation separate for use in the input to the task-specific RNN. Here is the `nn.Module`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ElmoRNNClassifierModel(TorchRNNClassifierModel):\n",
    "    def __init__(self, options_file, weights_file, hidden_dim, output_dim, bidirectional, device):     \n",
    "        super().__init__(\n",
    "            vocab_size=0,\n",
    "            embed_dim=1024, # self.elmo.get_output_dim()\n",
    "            use_embedding=False,\n",
    "            embedding=None,\n",
    "            hidden_dim=hidden_dim,\n",
    "            output_dim=output_dim,\n",
    "            bidirectional=bidirectional,\n",
    "            device=device)\n",
    "        self.options_file = options_file\n",
    "        self.weights_file = weights_file\n",
    "        self.elmo = Elmo(\n",
    "            self.options_file, \n",
    "            self.weights_file, \n",
    "            num_output_representations=2, \n",
    "            dropout=0)\n",
    "        \n",
    "    def forward(self, X, seq_lengths):\n",
    "        X = X.to(self.device, non_blocking=True)\n",
    "        result = self.elmo(X)\n",
    "        X = result['elmo_representations'][-1]\n",
    "        state = self.rnn_forward(X, seq_lengths, self.rnn)\n",
    "        logits = self.classifier_layer(state)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here is the subclass of `TorchRNNClassifier` that lets us take advantage of all the optimization and prediction methods of that class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ElmoRNNClassifier(TorchRNNClassifier):\n",
    "    def __init__(self, options_file, weights_file, *args, **kwargs):\n",
    "        self.options_file = options_file\n",
    "        self.weights_file = weights_file\n",
    "        vocab = []\n",
    "        super().__init__(vocab, *args, use_embedding=False, embedding=None, **kwargs)\n",
    "        \n",
    "    def build_graph(self):\n",
    "        \"\"\"This method is used by `fit`. We override it here to use our\n",
    "        new ELMo-based graph.\n",
    "        \n",
    "        \"\"\"\n",
    "        elmo = ElmoRNNClassifierModel(           \n",
    "            options_file=self.options_file,\n",
    "            weights_file=self.weights_file,\n",
    "            hidden_dim=self.hidden_dim,\n",
    "            output_dim=self.n_classes_,\n",
    "            bidirectional=self.bidirectional,\n",
    "            device=self.device)        \n",
    "        elmo.train()\n",
    "        return elmo\n",
    "    \n",
    "    def _prepare_dataset(self, X):\n",
    "        # Somewhat awkwardly get the lengths (padded tokens are all 0s).\n",
    "        # Ideally, this function would first measure the lengths of the\n",
    "        # examples in X directly and then call `batch_to_ids(X)`. \n",
    "        # However, the super class `fit` method is currently presupposing\n",
    "        # too much about X to allow this. TODO: make that interface\n",
    "        # more flexible.\n",
    "        seq_lengths = [sum([1 for w in ex if w.sum() > 0]) for ex in X]\n",
    "        return X, torch.tensor(seq_lengths)\n",
    "    \n",
    "    @staticmethod\n",
    "    def encode(X):\n",
    "        return batch_to_ids(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally here is a self-contained illustration:\n",
    "\n",
    "A copy of the dataset in the format we want:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_train = list(sst.train_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "elmo_dev = list(sst.dev_reader(SST_HOME, class_func=sst.ternary_class_func))\n",
    "\n",
    "X_elmo_tree_train, y_elmo_train = zip(*elmo_train)\n",
    "\n",
    "X_elmo_tree_dev, y_elmo_dev = zip(*elmo_dev)\n",
    "\n",
    "X_elmo_toks_train = [tree.leaves() for tree in X_elmo_tree_train]\n",
    "\n",
    "X_elmo_toks_dev = [tree.leaves() for tree in X_elmo_tree_dev]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our fine-tuning model with parameters that are inspired by those given for the SST task in [the paper's supplementary materials](https://www.aclweb.org/anthology/attachments/N18-1202.Notes.pdf):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_fine_tune_mod = ElmoRNNClassifier(\n",
    "    options_file, \n",
    "    weights_file, \n",
    "    batch_size=16, \n",
    "    max_iter=10,  # More iters improves performance. How many did the ELMo team do?\n",
    "    eta=0.0001,\n",
    "    l2_strength=0.0001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Character-level encodings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_elmo_train = elmo_fine_tune_mod.encode(X_elmo_toks_train)\n",
    "\n",
    "X_elmo_dev = elmo_fine_tune_mod.encode(X_elmo_toks_dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train (This experiment takes about 11 minutes on a machine with an NVIDIA RTX 2080 Max-Q GPU):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 10 of 10; error is 229.4948596879846"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8h 23min 11s, sys: 26min 34s, total: 8h 49min 46s\n",
      "Wall time: 2h 33min 1s\n"
     ]
    }
   ],
   "source": [
    "%time _ = elmo_fine_tune_mod.fit(X_elmo_train, y_elmo_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assess:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "# When I was on a GPU machine, I had trouble predicting on the \n",
    "# full dev set all at once, so let's break it into small batches:\n",
    "increment = 10\n",
    "elmo_fine_tune_preds = []\n",
    "for i in range(0, len(X_elmo_dev), increment):\n",
    "    elmo_fine_tune_preds += elmo_fine_tune_mod.predict(X_elmo_dev[i: i+increment])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.801     0.670     0.730       512\n",
      "     neutral      0.148     0.362     0.211        94\n",
      "    positive      0.813     0.729     0.769       495\n",
      "\n",
      "    accuracy                          0.670      1101\n",
      "   macro avg      0.588     0.587     0.570      1101\n",
      "weighted avg      0.751     0.670     0.703      1101\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(elmo_fine_tune_preds, y_elmo_dev, digits=3))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}