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, Fall 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. [HfBertClassifier](#HfBertClassifier)\n",
    "    1. [HfBertClassifier experiment](#HfBertClassifier-experiment)\n",
    "1. [Using ELMo](#Using-ELMo)\n",
    "  1. [ELMo Allen NLP set-up](#ELMo-Allen-NLP-set-up)\n",
    "  1. [ELMo fine-tuning](#ELMo-fine-tuning)\n",
    "    1. [AllenNLP ELMo interfaces](#AllenNLP-ELMo-interfaces)\n",
    "    1. [ElmoRNNClassifier](#ElmoRNNClassifier)\n",
    "    1. [ElmoRNNClassifier experiment](#ElmoRNNClassifier-experiment)"
   ]
  },
  {
   "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 laptop with a single NVIDIA RTX 2080 GPU. "
   ]
  },
  {
   "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 TorchRNNModel\n",
    "from torch_rnn_classifier import TorchRNNClassifier\n",
    "from torch_rnn_classifier import 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": [
    "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, 2.5, and 2.11 of `transformers`. Try to get at least 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": [],
   "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": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3cc30d02a7d94499a5db796f860340c0"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "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 the shape of the tensor for 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 computes 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": [
    "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.binary_class_func))\n",
    "\n",
    "hf_dev = list(sst.dev_reader(SST_HOME, class_func=sst.binary_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 34min 55s, sys: 6.99 s, total: 35min 2s\n",
      "Wall time: 5min 52s\n"
     ]
    }
   ],
   "source": [
    "%time X_hf_train = [hugging_face_bert_classifier_phi(tree)\n",
    "                    for tree in X_hf_tree_train]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4min 13s, sys: 868 ms, total: 4min 13s\n",
      "Wall time: 42.6 s\n"
     ]
    }
   ],
   "source": [
    "%time X_hf_dev = [hugging_face_bert_classifier_phi(tree)\n",
    "                  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(\n",
    "    early_stopping=True,\n",
    "    hidden_dim=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 38. Validation score did not improve by tol=1e-05 for more than 10 epochs. Final error is 2.0183719843626022"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 19.2 s, sys: 496 ms, total: 19.7 s\n",
      "Wall time: 5.28 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.850     0.834     0.842       428\n",
      "    positive      0.843     0.858     0.850       444\n",
      "\n",
      "    accuracy                          0.846       872\n",
      "   macro avg      0.846     0.846     0.846       872\n",
      "weighted avg      0.846     0.846     0.846       872\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",
    "        hidden_dim=300,\n",
    "        early_stopping=True)\n",
    "    mod.fit(X, y)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 22. Validation score did not improve by tol=1e-05 for more than 10 epochs. Final error is 2.3845930993556976"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.839     0.839     0.839       428\n",
      "    positive      0.845     0.845     0.845       444\n",
      "\n",
      "    accuracy                          0.842       872\n",
      "   macro avg      0.842     0.842     0.842       872\n",
      "weighted avg      0.842     0.842     0.842       872\n",
      "\n",
      "CPU times: user 42min 27s, sys: 8.83 s, total: 42min 36s\n",
      "Wall time: 7min 8s\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.binary_class_func,\n",
    "    vectorize=False)  # Pass in the BERT reps 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",
    "        early_stopping=True,\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": [
      "Stopping after epoch 25. Validation score did not improve by tol=1e-05 for more than 10 epochs. Final error is 0.9420086964964867"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.849     0.853     0.851       428\n",
      "    positive      0.857     0.854     0.856       444\n",
      "\n",
      "    accuracy                          0.853       872\n",
      "   macro avg      0.853     0.853     0.853       872\n",
      "weighted avg      0.853     0.853     0.853       872\n",
      "\n",
      "CPU times: user 44min, sys: 31.2 s, total: 44min 31s\n",
      "Wall time: 7min 37s\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.binary_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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### HfBertClassifier\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. Here we define a very simple fine-tuning set-up in which some layers built on top of the output corresponding to `[CLS]` are used as the basis for the SST 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.bert.train()\n",
    "        self.hidden_dim = self.bert.embeddings.word_embeddings.embedding_dim\n",
    "        # The only new parameters -- the classifier:\n",
    "        self.classifier_layer = nn.Linear(\n",
    "            self.hidden_dim, self.n_classes)\n",
    "\n",
    "    def forward(self, indices, mask):\n",
    "        final_hidden_states, cls_output = self.bert(\n",
    "            indices, attention_mask=mask)\n",
    "        return self.classifier_layer(cls_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, `self.bert` does the heavy-lifting: it reads in all the pretrained BERT parameters, and I've specified `self.bert.train()` just to make sure that these parameters can be updated during our training process. \n",
    "\n",
    "In `forward`, `self.bert` is used to process inputs, and then `cls_output` is fed into `self.classifier_layer`. Hugging Face has already added a layer on top of the actual output for `[CLS]`, so we can specify the model as\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "[h_{1}, \\ldots, h_{n}] &= \\text{BERT}([x_{1}, \\ldots, x_{n}]) \\\\\n",
    "h &= \\tanh(h_{1}W_{hh} + b_{h}) \\\\\n",
    "y &= \\textbf{softmax}(hW_{hy} + b_{y})\n",
    "\\end{align}$$\n",
    "\n",
    "for a tokenized input sequence $[x_{1}, \\ldots, x_{n}]$. \n",
    "\n",
    "The Hugging Face documentation somewhat amusingly says, of `cls_output`,\n",
    "\n",
    "> This output is usually _not_ a good summary of the semantic content of the input, you're often better with averaging or pooling the sequence of hidden-states for the whole input sequence.\n",
    "\n",
    "which is entirely reasonable, but it will require more resources, so we'll do the simpler thing here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the training and prediction interface, we can subclass `TorchShallowNeuralClassifier` so that we don't have to write any of our own data-handling, training, or prediction code. The central changes are using `HfBertClassifierModel` in `build_graph` and processing the data with `batch_encode_plus`."
   ]
  },
  {
   "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",
    "        self.params += ['weights_name']\n",
    "\n",
    "    def build_graph(self):\n",
    "        return HfBertClassifierModel(self.n_classes_, self.weights_name)\n",
    "\n",
    "    def build_dataset(self, X, y=None):\n",
    "        data = self.tokenizer.batch_encode_plus(\n",
    "            X,\n",
    "            max_length=None,\n",
    "            add_special_tokens=True,\n",
    "            pad_to_max_length=True,\n",
    "            return_attention_mask=True)\n",
    "        indices = torch.tensor(data['input_ids'])\n",
    "        mask = torch.tensor(data['attention_mask'])\n",
    "        if y is None:\n",
    "            dataset = torch.utils.data.TensorDataset(indices, mask)\n",
    "        else:\n",
    "            self.classes_ = sorted(set(y))\n",
    "            self.n_classes_ = len(self.classes_)\n",
    "            class2index = dict(zip(self.classes_, range(self.n_classes_)))\n",
    "            y = [class2index[label] for label in y]\n",
    "            y = torch.tensor(y)\n",
    "            dataset = torch.utils.data.TensorDataset(indices, mask, y)\n",
    "        return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### HfBertClassifier experiment\n",
    "\n",
    "That's it! Let's see how we do on the SST binary, root-only problem. Because fine-tuning is expensive, we'll conduct a modest hyperparameter search and run the model for just one epoch per setting evaluation, as we did when [assessing NLI models](nli_02_models.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bert_fine_tune_phi(tree):\n",
    "    return \" \".join(tree.leaves())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_hf_bert_classifier_with_hyperparameter_search(X, y):\n",
    "    basemod = HfBertClassifier(\n",
    "        weights_name='bert-base-cased',\n",
    "        batch_size=8,  # Small batches to avoid memory overload.\n",
    "        max_iter=1,  # We'll search based on 1 iteration for efficiency.\n",
    "        n_iter_no_change=5,   # Early-stopping params are for the\n",
    "        early_stopping=True)  # final evaluation.\n",
    "\n",
    "    param_grid = {\n",
    "        'gradient_accumulation_steps': [1, 4, 8],\n",
    "        'eta': [0.00005, 0.0001, 0.001],\n",
    "        'hidden_dim': [100, 200, 300]}\n",
    "\n",
    "    bestmod = utils.fit_classifier_with_hyperparameter_search(\n",
    "        X, y, basemod, cv=3, param_grid=param_grid)\n",
    "\n",
    "    return bestmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 1 of 1; error is 38.469218485057354"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best params: {'eta': 5e-05, 'gradient_accumulation_steps': 8, 'hidden_dim': 300}\n",
      "Best score: 0.892\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.953     0.808     0.875       428\n",
      "    positive      0.839     0.962     0.896       444\n",
      "\n",
      "    accuracy                          0.886       872\n",
      "   macro avg      0.896     0.885     0.885       872\n",
      "weighted avg      0.895     0.886     0.886       872\n",
      "\n",
      "CPU times: user 1h 31min 39s, sys: 4min 6s, total: 1h 35min 46s\n",
      "Wall time: 1h 36min 43s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "hf_bert_classifier_xval = sst.experiment(\n",
    "    SST_HOME,\n",
    "    bert_fine_tune_phi,\n",
    "    fit_hf_bert_classifier_with_hyperparameter_search,\n",
    "    train_reader=sst.train_reader,\n",
    "    assess_reader=sst.dev_reader,\n",
    "    class_func=sst.binary_class_func,\n",
    "    vectorize=False)  # Pass in the BERT hidden state directly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now on to the final test-set evaluation, using the best model from above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_hf_bert_classifier = hf_bert_classifier_xval['model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove the rest of the experiment results to clear out some memory:\n",
    "del hf_bert_classifier_xval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_optimized_hf_bert_classifier(X, y):\n",
    "    optimized_hf_bert_classifier.max_iter = 1000\n",
    "    optimized_hf_bert_classifier.fit(X, y)\n",
    "    return optimized_hf_bert_classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 10. Validation score did not improve by tol=1e-05 for more than 5 epochs. Final error is 0.964973742607981"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.890     0.928     0.908       912\n",
      "    positive      0.924     0.884     0.904       909\n",
      "\n",
      "    accuracy                          0.906      1821\n",
      "   macro avg      0.907     0.906     0.906      1821\n",
      "weighted avg      0.907     0.906     0.906      1821\n",
      "\n",
      "CPU times: user 13min 12s, sys: 17.4 s, total: 13min 29s\n",
      "Wall time: 13min 29s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "_ = sst.experiment(\n",
    "    SST_HOME,\n",
    "    bert_fine_tune_phi,\n",
    "    fit_optimized_hf_bert_classifier,\n",
    "    train_reader=(sst.train_reader, sst.dev_reader),\n",
    "    assess_reader=sst.test_reader,\n",
    "    class_func=sst.binary_class_func,\n",
    "    vectorize=False)  # Pass in the BERT hidden state directly!"
   ]
  },
  {
   "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, 0.9.0, and 1.0.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": [],
   "source": [
    "from allennlp.modules.elmo import Elmo, batch_to_ids\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://docs.allennlp.org/master/api/modules/elmo/"
   ]
  },
  {
   "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`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AllenNLP ELMo interfaces\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": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo = Elmo(options_file, weights_file, num_output_representations=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model expects tokenized inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "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": 46,
   "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": 52,
     "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": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "elmo_embeddings = elmo(elmo_character_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`elmo_embeddings` is a dict. The value of the key `'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. I specified `num_output_representations=1` when initializing `elmo` above, so we get a list of length 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[[-0.7944,  0.0000,  0.0000,  ..., -0.0000, -0.0000, -1.9283],\n",
       "          [-0.0000,  1.1850,  0.6947,  ...,  0.0000, -0.9297, -0.2358],\n",
       "          [ 0.0000,  0.5358,  0.7767,  ..., -0.6500, -0.0777, -0.4875],\n",
       "          [-0.0000, -0.4965, -0.0000,  ..., -0.1605,  0.0000,  0.2256]],\n",
       " \n",
       "         [[ 0.4553, -0.0000,  0.0000,  ..., -0.0000, -0.0000, -0.5380],\n",
       "          [ 0.0000, -0.5309, -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>)]"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elmo_embeddings['elmo_representations']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ElmoRNNClassifier\n",
    "\n",
    "The above are the representations we will be fine-tuning. There are many ways to cdo this. 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` built on `TorchRNNClassifierModel`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ElmoRNNClassifierModel(TorchRNNClassifierModel):\n",
    "    def __init__(self,\n",
    "            options_file,\n",
    "            weights_file,\n",
    "            rnn,\n",
    "            output_dim,\n",
    "            classifier_activation):\n",
    "        super().__init__(rnn, output_dim, classifier_activation)\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=1,\n",
    "            dropout=0)\n",
    "        self.elmo.train()\n",
    "\n",
    "    def forward(self, X, seq_lengths):\n",
    "        result = self.elmo(X)\n",
    "        X = result['elmo_representations'][-1]\n",
    "        outputs, state  = self.rnn(X, seq_lengths)\n",
    "        state = self.get_batch_final_states(state)\n",
    "        if self.rnn.bidirectional:\n",
    "            state = torch.cat((state[0], state[1]), dim=1)\n",
    "        h = self.classifier_activation(self.hidden_layer(state))\n",
    "        logits = self.classifier_layer(h)\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": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ElmoRNNClassifier(TorchRNNClassifier):\n",
    "    def __init__(self, vocab, options_file, weights_file, **model_kwargs):\n",
    "        self.options_file = options_file\n",
    "        self.weights_file = weights_file\n",
    "        # Values determined by using ELMo:\n",
    "        model_kwargs['use_embedding'] = False\n",
    "        model_kwargs['embedding'] = None\n",
    "        model_kwargs['embed_dim'] = 1024\n",
    "        super().__init__(vocab, **model_kwargs)\n",
    "        self.params += ['options_file', 'weights_file']\n",
    "\n",
    "    def build_graph(self):\n",
    "\n",
    "        # The RNN is setup just as in a regular `TorchRNNClassifier`:\n",
    "        rnn = TorchRNNModel(\n",
    "            vocab_size=len(self.vocab),\n",
    "            embedding=self.embedding,\n",
    "            use_embedding=self.use_embedding,\n",
    "            embed_dim=self.embed_dim,\n",
    "            rnn_cell_class=self.rnn_cell_class,\n",
    "            hidden_dim=self.hidden_dim,\n",
    "            bidirectional=self.bidirectional,\n",
    "            freeze_embedding=self.freeze_embedding)\n",
    "\n",
    "        # The Classifier layer uses our new `ElmoRNNClassifierModel`:\n",
    "        model = ElmoRNNClassifierModel(\n",
    "            options_file=self.options_file,\n",
    "            weights_file=self.weights_file,\n",
    "            rnn=rnn,\n",
    "            output_dim=self.n_classes_,\n",
    "            classifier_activation=self.classifier_activation)\n",
    "\n",
    "        return model\n",
    "\n",
    "    def build_dataset(self, X, y=None):\n",
    "        seq_lengths = [len(ex) for ex in X]\n",
    "        seq_lengths = torch.tensor(seq_lengths)\n",
    "        X = batch_to_ids(X)\n",
    "        if y is None:\n",
    "            return torch.utils.data.TensorDataset(X, seq_lengths)\n",
    "        else:\n",
    "            self.classes_ = sorted(set(y))\n",
    "            self.n_classes_ = len(self.classes_)\n",
    "            class2index = dict(zip(self.classes_, range(self.n_classes_)))\n",
    "            y = [class2index[label] for label in y]\n",
    "            y = torch.tensor(y)\n",
    "            return torch.utils.data.TensorDataset(X, seq_lengths, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ElmoRNNClassifier experiment\n",
    "\n",
    "And finally here is a self-contained evaluation involving a modest hyperparameter search:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elmo_fine_tune_phi(tree):\n",
    "    return tree.leaves()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_elmo_rnn_with_hyperparameter_search(X, y):\n",
    "    basemod = ElmoRNNClassifier(\n",
    "        vocab=[],\n",
    "        options_file=options_file,\n",
    "        weights_file=weights_file,\n",
    "        batch_size=8,  # Kept small so that we can explore large networks.\n",
    "        max_iter=1,  # We'll search based on 1 iteration for efficiency.\n",
    "        n_iter_no_change=5,   # Early-stopping params are for the\n",
    "        early_stopping=True)  # final evalution.\n",
    "\n",
    "    param_grid = {\n",
    "        'gradient_accumulation_steps': [1, 4, 8],\n",
    "        'eta': [0.001, 0.01, 0.05],\n",
    "        'hidden_dim': [50, 100, 200]}\n",
    "\n",
    "    bestmod = utils.fit_classifier_with_hyperparameter_search(\n",
    "        X, y, basemod, cv=3, param_grid=param_grid)\n",
    "\n",
    "    return bestmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finished epoch 1 of 1; error is 330.32408917695284"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best params: {'eta': 0.001, 'gradient_accumulation_steps': 1, 'hidden_dim': 200}\n",
      "Best score: 0.856\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.848     0.848     0.848       428\n",
      "    positive      0.854     0.854     0.854       444\n",
      "\n",
      "    accuracy                          0.851       872\n",
      "   macro avg      0.851     0.851     0.851       872\n",
      "weighted avg      0.851     0.851     0.851       872\n",
      "\n",
      "CPU times: user 2h 4min 55s, sys: 4min 16s, total: 2h 9min 12s\n",
      "Wall time: 2h 1min 35s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "elmo_rnn_xval = sst.experiment(\n",
    "    SST_HOME,\n",
    "    elmo_fine_tune_phi,\n",
    "    fit_elmo_rnn_with_hyperparameter_search,\n",
    "    train_reader=sst.train_reader,\n",
    "    assess_reader=sst.dev_reader,\n",
    "    class_func=sst.binary_class_func,\n",
    "    vectorize=False)  # Pass in the ELMo reps directly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we move to the test-set evaluation using the best model we found:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_elmo_rnn = elmo_rnn_xval['model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove the unneeded experimental data to save memory:\n",
    "del elmo_rnn_xval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_optimized_elmo_rnn(X, y):\n",
    "    optimized_elmo_rnn.max_iter = 20\n",
    "    optimized_elmo_rnn.fit(X, y)\n",
    "    return optimized_elmo_rnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 10. Validation score did not improve by tol=1e-05 for more than 5 epochs. Final error is 27.70847176760435"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative      0.891     0.880     0.886       912\n",
      "    positive      0.882     0.892     0.887       909\n",
      "\n",
      "    accuracy                          0.886      1821\n",
      "   macro avg      0.886     0.886     0.886      1821\n",
      "weighted avg      0.886     0.886     0.886      1821\n",
      "\n",
      "CPU times: user 15min 31s, sys: 34.3 s, total: 16min 6s\n",
      "Wall time: 16min\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "_ = sst.experiment(\n",
    "    SST_HOME,\n",
    "    elmo_fine_tune_phi,\n",
    "    fit_optimized_elmo_rnn,\n",
    "    train_reader=(sst.train_reader, sst.dev_reader),\n",
    "    assess_reader=sst.test_reader,\n",
    "    class_func=sst.binary_class_func,\n",
    "    vectorize=False)  # Pass in the BERT hidden state directly!"
   ]
  }
 ],
 "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
}