{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coreference Resolution\n",
    "\n",
    "In this problem set, you will venture into the challenging NLP task of **coreference resolution**. You will:\n",
    "\n",
    "* Implement a simple rule-based system that achieve results which are surprisingly difficult to beat.\n",
    "* Get acquainted with the trickiness of evaluating coref systems, and the current solutions in the field.\n",
    "* Experiment with two neural approaches for coref to be implemented in [PyTorch](pytorch.org):\n",
    "  * A feedforward network that only looks at boolean mention-pair features\n",
    "  * A fully-neural architecture with embeddings all the way down\n",
    "* Get a glimpse at **domain adaptation** in the wild, by trying to run a system trained on news against a narrative corpus and vice-versa.\n",
    "\n",
    "\n",
    "# 0. Setup\n",
    "\n",
    "In order to develop this assignment, you will need [python 3.6](https://www.python.org/downloads/) and the following libraries. Most if not all of these are part of [anaconda](https://www.continuum.io/downloads), so a good starting point would be to install that.\n",
    "\n",
    "* [jupyter](http://jupyter.readthedocs.org/en/latest/install.html)\n",
    "* [numpy](https://docs.scipy.org/doc/numpy/user/install.html)\n",
    "* [nosetests](https://nose.readthedocs.org/en/latest/)\n",
    "* [pytorch](https://pytorch.org/)\n",
    "* [scikit-learn](http://scikit-learn.org/)\n",
    "* [xmltodict](https://github.com/martinblech/xmltodict#ok-how-do-i-get-it)\n",
    "\n",
    "Here is some help on installing packages in python: https://packaging.python.org/installing/. You can use pip --user to install locally without sudo.\n",
    "\n",
    "## About this assignment\n",
    "\n",
    "* This is a Jupyter notebook. You can execute cell blocks by pressing control+enter, or execute and move to the next using shift+enter, or execute and open a new cell below using alt+enter.\n",
    "* Most of your coding will be in the python source files in the directory `gtnlplib`.\n",
    "* The directory `tests` contains unit tests that will be used to grade your assignment, using `nosetests`. You should run them as you work on the assignment to see that you're on the right track. You are free to look at their source code, if that helps -- though most of the relevant code is also here in this notebook. Learn more about running unit tests [here](http://pythontesting.net/framework/nose/nose-introduction/).\n",
    "* You may want to add more tests, but that is completely optional.\n",
    "* **To submit this assignment, run the script `make-submission.sh`, and submit the tarball `pset4-submission.tgz` on Canvas.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import nose\n",
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "import sklearn\n",
    "import xmltodict\n",
    "from collections import Counter, defaultdict\n",
    "from nltk.tag import pos_tag\n",
    "\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "# no more reload cells!\n",
    "# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "My library versions\n",
      "numpy: 1.14.0\n",
      "xmltodict: 0.11.0\n",
      "nose: 1.3.7\n",
      "sklearn: 0.19.1\n",
      "torch: 0.3.0b0+591e73e\n"
     ]
    }
   ],
   "source": [
    "print('My library versions')\n",
    "\n",
    "print('numpy: {}'.format(np.__version__))\n",
    "print('xmltodict: {}'.format(xmltodict.__version__))\n",
    "print('nose: {}'.format(nose.__version__))\n",
    "print('sklearn: {}'.format(sklearn.__version__))\n",
    "print('torch: {}'.format(torch.__version__)) # this TA uses a Windows machine, so YMMV"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test whether your libraries are the right version, run:\n",
    "\n",
    "`nosetests tests/test_environment.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ".\n",
      "----------------------------------------------------------------------\n",
      "Ran 1 test in 0.000s\n",
      "\n",
      "OK\n"
     ]
    }
   ],
   "source": [
    "# use ! to run shell commands in notebook\n",
    "! nosetests tests/test_environment.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from gtnlplib import coref, coref_rules, coref_features, coref_learning, neural_net, utils\n",
    "\n",
    "# constants for notebook use\n",
    "ETA_0 = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1: Exploring the data\n",
    "\n",
    "The core data is in the form of \"markables\", or \"referring expressions\", which refer to token sequences that can participate in coreference relations.\n",
    "\n",
    "Each markable is a namedtuple with five elements:\n",
    "- ```string```, which is a list of tokens\n",
    "- ```entity```, which defines the ground truth assignments\n",
    "- ```start_token```, the index of the first token in the markable with respect to the entire document\n",
    "- ```end_token```, one plus the index of the last token in the markable\n",
    "- ```tags```, POS tags corresponding to the tokens in ```string``` which will remain NULL for now\n",
    "\n",
    "The ```read_data``` function also returns a list of tokens.\n",
    "You can use this to incorporate the linguistic context around each markable.\n",
    "\n",
    "### Loading the dataset\n",
    "\n",
    "For most of this problem set, we will explore a dataset of articles from the Wall Street Journal (WSJ) extracted and annotated from the Penn Treebank (PTB)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dv_dir = os.path.join('data','wsj','dev')\n",
    "tr_dir = os.path.join('data','wsj','train')\n",
    "te_dir = os.path.join('data','wsj','test') # all markables here are annotated as the same entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "markables, words = coref.read_data('06_wsj_0051.sty',basedir=tr_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Markable object: Markable(string=['Fujitsu', 'Ltd.'], entity='set_3082', start_token=0, end_token=2, tags=['NULL', 'NULL'])\n",
      "Words for markable extracted from text: ['Fujitsu', 'Ltd.']\n"
     ]
    }
   ],
   "source": [
    "print('Markable object:', markables[0])\n",
    "print('Words for markable extracted from text:', words[markables[0].start_token:markables[0].end_token])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 1.1**: Write a function that returns all the markable **strings** associated with a given entity. Specifically, fill in the function `get_markables_for_entity()` in `coref.py`.\n",
    "(4650: 1 point; 7650: 0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_get_markables_d1_1()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Fujitsu',\n",
       " 'Fujitsu',\n",
       " 'Fujitsu',\n",
       " 'Fujitsu',\n",
       " 'Fujitsu',\n",
       " 'Fujitsu',\n",
       " 'Fujitsu',\n",
       " \"Fujitsu , Japan 's No. 1 computer maker\",\n",
       " 'Fujitsu Ltd.',\n",
       " 'It',\n",
       " 'The company',\n",
       " 'The company',\n",
       " 'We',\n",
       " 'his company',\n",
       " 'his company',\n",
       " 'it',\n",
       " 'it',\n",
       " 'it',\n",
       " 'it',\n",
       " 'it',\n",
       " 'it',\n",
       " 'its',\n",
       " 'its',\n",
       " 'the company']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sorted(coref.get_markables_for_entity(markables,'set_3082'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 1.2** Write a function that takes as input a string, and returns a list of distances to the most recent ground truth antecedent for every time the (case-insensitive) input string appears. For example, if the input is \"they\", it should make a list with one element for each time the word \"they\" appears in the list of markables. Each element should be the distance of the word \"they\" to the nearest previous mention of the entity that \"they\" references.\n",
    "\n",
    "Fill in the function `get_distances()` in `coref.py`. If the input string is not anaphoric, the distance should be zero. Note that input strings may contain spaces. You may use any other function in `coref.py` to help you.\n",
    "(4650: 1 point; 7650: 0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_get_antecedents_d1_2()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 1, 1, 2, 2]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coref.get_distances(markables,'they')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's compare the typical distances for various mention types.\n",
    "\n",
    "You can see the most frequent mention types by using the `Counter` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('it', 9), ('Fujitsu', 7), ('Japan', 5), ('they', 5), ('NEC', 4)]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Counter([' '.join(markable.string) for markable in markables]).most_common(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[15, 8, 49, 12, 7, 4, 2]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coref.get_distances(markables, 'Fujitsu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[4, 4, 6]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coref.get_distances(markables, 'the company')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 2, 2, 1, 6, 1, 6, 2, 6, 1]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coref.get_distances(markables, 'it') # there are 10 because our counter was case-sensitive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Rule-based coreference resolution\n",
    "\n",
    "We have written a simple coreference classifier, which predicts that each markable is linked to the most recent antecedent which is an exact string match.\n",
    "\n",
    "The code block below applies this method to the dev set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "exact_matcher = coref_rules.make_resolver(coref_rules.exact_match)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code above has two pieces:\n",
    "\n",
    "- ```coref_rules.exact_match()``` is a function that takes two markables, and returns `True` iff they are an exact (case-insensitive) string match\n",
    "- ```make_resolver()``` is a function that takes a matching function, and returns a function that computes an antecedent list for a list of markables.\n",
    "\n",
    "Let's run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]\n"
     ]
    }
   ],
   "source": [
    "ant_exact = exact_matcher(markables)\n",
    "print(ant_exact[:20])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output is a list of antecedent numbers, $c_i$. \n",
    "When $c_i = i$, the markable $i$ has no antecedent: it is the first mention of its entity. In this case, all first 20 mentions are new and don't have antecedents. You can try and modify the cell to look further down the list to see if there are ever actual matches, which we know should occur due to the output of cell 5.\n",
    "\n",
    "We can test whether these predictions are correct by comparing against the key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ant_true = coref.get_true_antecedents(markables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "correct: 128\taccuracy: 0.660\n"
     ]
    }
   ],
   "source": [
    "num_correct = sum([c_true == c_predict for c_true, c_predict in zip(ant_true, ant_exact)])\n",
    "acc = num_correct / len(markables)\n",
    "print(f'correct: {num_correct}\\taccuracy: {acc:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "Coreference can be evaluated in terms of recall, precision, and F-measure. Here is how we will define these terms:\n",
    "\n",
    "- **True positive**: The system predicts $\\hat{c}_i < i$, and $\\hat{c}_i$ and $i$ are references to the same entity.\n",
    "- **False positive**: The system predicts $\\hat{c}_i < i$, but $\\hat{c}_i$ and $i$ are not references to the same entity.\n",
    "- **False negative**: There exists some $c_i < i$ such that $c_i$ and $i$ are references to the same entity, but the system predicts either $\\hat{c}_i = i$, or some $\\hat{c}_i$ which is not really a reference to the same entity that $i$ references.\n",
    "- Recall = $\\frac{tp}{tp + fn}$\n",
    "- Precision = $\\frac{tp}{tp + fp}$\n",
    "- F-measure = $\\frac{2RP}{R+P}$\n",
    "\n",
    "A couple of things to notice here:\n",
    "\n",
    "- There is no reward for correctly identifying a markable as non-anaphoric (not having any antecedent), but you do avoid committing a false positive by doing this.\n",
    "- You cannot compute the evaluation by directly matching the predicted antecedents to the true antecedents. Suppose the truth is $a \\leftarrow b, b \\leftarrow c$, but the system predicts $a \\leftarrow b, a \\leftarrow c$: the system should receive two true positives, since $a$ and $c$ are references to the same entity in the ground truth.\n",
    "\n",
    "**Deliverable 2.1** Implement `get_tp()`, `get_fp()`, and `get_fn()` in `coref.py`. You will want to use the function `coref.get_entities()`.\n",
    "(1 point)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_recall_d2_1(), test_precision_d2_1(), test_fmeasure_d2_1()`\n",
    "\n",
    "**NOTE!** You **must** successfully complete this deliverable. Otherwise, some of the unit tests won't work and you won't be able to complete the rest of the assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5231\t0.4096\t0.7234\n"
     ]
    }
   ],
   "source": [
    "f,r,p = coref.evaluate_f(exact_matcher, markables)\n",
    "print(f'{f:.4f}\\t{r:.4f}\\t{p:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "all_markables, all_words = coref.read_dataset(tr_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5452\tR: 0.4130\tP:0.8018\n"
     ]
    }
   ],
   "source": [
    "coref.eval_on_dataset(exact_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before optimizing on this simple F-measure (sometimes called F1) and its components, one should be aware that in the real world coreference is evaluated over three other metrics, namely $B^3$, **`CEAF`**, and **`MUC`**. We have implemented them for you, and will check our performance on them from time to time. You can read more about them [here](http://www.anthology.aclweb.org/W/W10/W10-4305.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def coref_metrics(matcher, dataset):\n",
    "    ants = [matcher(m) for m in dataset]\n",
    "    b3, ceaf, muc = coref.evaluate_bcm(dataset, ants)\n",
    "    avg_f1 = np.average([b3, ceaf, muc])\n",
    "    print(f'B-Cubed: {b3:.4f}\\tCEAF: {ceaf:.4f}\\tMUC: {muc:.4f}\\tAverage: {avg_f1:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5756\tCEAF: 0.4551\tMUC: 0.5605\tAverage: 0.5304\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(exact_matcher, all_markables); #\"Average\" is the commonly used main metric in state-of-the-art systems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reasons for having multiple measures for coref evaluation is manyfold. One of them, discussed in the paper linked to above, has to do with the pre-resolution task of *identifying markables*. Since we're only working with pre-extracted markables, that needn't worry us.\n",
    "\n",
    "The other reason is that different perspectives on coreference matching are equally plausible - we can focus on single correct predictions, or finding the correct clusters for each entity , for example. Each of these is \"gameable\" by different trivial classifiers.\n",
    "\n",
    "**Deliverable 2.2** To witness this problem, you will implement `coref_rules.singleton_matcher()`, which produces an assignment where each markable has its own entity, and `coref_rules.full_cluster_matcher()`, which assigns all markables to the same entity. Running the metrics against them will demonstrate the problem.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_singleton_matcher_d2_2(), test_full_cluster_matcher_d2_2()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.3411\tCEAF: 0.0012\tMUC: 0.0000\tAverage: 0.1141\n"
     ]
    }
   ],
   "source": [
    "singleton_resolver = coref_rules.make_resolver(coref_rules.singleton_matcher)\n",
    "coref_metrics(singleton_resolver, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MUC has an inherent problem evaluating singleton entities. $B^3$, on the other hand, is extra-generous with them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.0739\tCEAF: 0.0304\tMUC: 0.5552\tAverage: 0.2199\n"
     ]
    }
   ],
   "source": [
    "full_cluster_resolver = coref_rules.make_resolver(coref_rules.full_cluster_matcher)\n",
    "coref_metrics(full_cluster_resolver, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case MUC, which is focused on detecting incompatible clusters, is fairly comfortable with the fact that there's only one predicted cluster.\n",
    "CEAF, a metric which gives low precision very easily, is difficult to score high on."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Increasing precision\n",
    "\n",
    "The `exact_match()` function matches everything, including pronouns. This can lead to mistakes:\n",
    "\n",
    "\"Umashanthi ate pizza until she was full. Parvati kept eating until she had a stomach ache.\"\n",
    "\n",
    "In this example, both pronouns likely refer to the names that immediately precede them, and not to each other.\n",
    "\n",
    "**Deliverable 2.3** The file `coref_rules.py` contains the signature for a function `exact_match_no_pronoun()`, which solves this problem by only predicting matches between markables that are not pronouns. Implement and test this function. For now, you may use the list of pronouns provided in the code file `coref_rules.py`.\n",
    "(0.25 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_match_nopro_d2_3(), tests\\test_coref.py:test_match_nopro_f1_d2_3()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "no_pro_matcher = coref_rules.make_resolver(coref_rules.exact_match_no_pronouns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.4551\tR: 0.3028\tP:0.9158\n"
     ]
    }
   ],
   "source": [
    "f,r,p = coref.eval_on_dataset(no_pro_matcher,all_markables);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5678\tCEAF: 0.4269\tMUC: 0.4568\tAverage: 0.4839\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(no_pro_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Precision has increased, but recall decreased, dragging down the overall F-measure as well as our favorite metrics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Increasing recall\n",
    "\n",
    "Our current matcher is very conservative. Let's try to increase recall. One solution is match on the **head word** of each markable. \n",
    "\n",
    "As you know, in a CFG parse, the head word is defined by a set of rules: for example, the head of a determiner-noun construction is the noun. In a dependency parse, the head word would be the root of the subtree governing the markable span. But this assumes that the markables correspond to syntactic constituents or dependency subtrees. This is not guaranteed to be true - particularly when there are parsing errors.\n",
    "\n",
    "**Deliverable 2.4** Let's start with a much simpler head-finding heuristic: simply select the *last word* in the markable. This handles many cases - but as we will see, not all. To do this, implement the function `match_last_token()` in ```coref_rules.py```. This function should match all cases where the final tokens match.\n",
    "(0.25 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_match_last_tok_d2_4()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "last_tok_matcher = coref_rules.make_resolver(coref_rules.match_last_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.3994\tR: 0.4385\tP:0.3666\n"
     ]
    }
   ],
   "source": [
    "coref.eval_on_dataset(last_tok_matcher,all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall is up, but precision is back down. To try to increase precision, let's add one more rule: two markables cannot coref if their spans overlap. This can happen with nested mentions, such as \"(the president (of the United States))\". Under our last-token rule, these two mentions would co-refer, but logically, overlapping markables cannot refer to the same entity. \n",
    "\n",
    "**Deliverable 2.5** Fill in the function `match_last_token_no_overlap()`, which should match any two markables that share the same last token, unless their spans overlap. Use the `start_token` and `end_token` members of each markable to determine whether they overlap. the final tokens match.\n",
    "(0.25 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_match_no_overlap_f1_d2_5()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.4911\tR: 0.4965\tP:0.4858\n"
     ]
    }
   ],
   "source": [
    "mltno_matcher = coref_rules.make_resolver(coref_rules.match_last_token_no_overlap)\n",
    "coref.eval_on_dataset(mltno_matcher,all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both recall and precision increase. Why would recall increase? The restriction does not create any new coreference links, but it changes some incorrect links to correct links. This increases the number of true positives and reduces the number of false negatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5571\tCEAF: 0.4610\tMUC: 0.5473\tAverage: 0.5218\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(mltno_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Almost back to the results from the exact matcher."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Error analysis\n",
    "\n",
    "To see whether we can do even better, let's try some error analysis on a specific file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# predicted antecedent series\n",
    "markables_17, _ = coref.read_data('17_wsj_0072.sty',basedir=tr_dir)\n",
    "ant = coref_rules.make_resolver(coref_rules.match_last_token_no_overlap)(markables_17)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# let's look at large entities\n",
    "m2e, e2m = coref.markables_to_entities(markables_17,ant)\n",
    "big_entities = [ent for ent, vals in e2m.items() if len(vals) > 10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entity 8: 11 mentions\n",
      "['Fed', 'Kansas City Fed', 'the Fed', 'the Fed', 'The Fed', 'The report from the Fed', 'the Fed', 'The Philadelphia Fed', 'Fed', 'Fed', 'regional Fed']\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for entity in big_entities:\n",
    "    print(f'Entity {entity}: {len(e2m[entity])} mentions')\n",
    "    print([' '.join(markables_17[idx].string) for idx in e2m[entity]])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Incorporating parts of speech\n",
    "\n",
    "One clear mistake is that we are matching ''Kansas City Fed'' to ''The Philadelphia Fed'' and other ''Fed''s. The last token heuristic is the culprit: in this case, the first token is a key disambiguator. Let's try a more syntactically-motivated approach. \n",
    "\n",
    "Instead of matching the last token (low precision) or matching on all tokens (low recall), let's try matching on all *content* words. Let's start by including only the following grammatical categories:\n",
    "\n",
    "- Nouns (proper, common, singular, plural)\n",
    "- Pronouns (including possessive)\n",
    "- Adjectives (including comparative and superlative)\n",
    "- Cardinal numbers\n",
    "\n",
    "To get these categories, we can call `read_dataset()` again with the optional `tagger` argument, a part of speech tagger. We'll use NLTK for this project, which has a structured perceptron tagger on the [PTB tagset](https://www.ling.upenn.edu/courses/Fall_2003/ling001/penn_treebank_pos.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "all_markables, _ = coref.read_dataset(tr_dir, tagger=pos_tag)\n",
    "all_markables_dev, all_words_dev = coref.read_dataset(dv_dir, tagger=pos_tag)\n",
    "all_markables_te, all_words_test = coref.read_dataset(te_dir, tagger=pos_tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Markable(string=['a', 'trade', 'deficit', 'of', '$', '101', 'million'], entity='set_972', start_token=3, end_token=10, tags=['DT', 'NN', 'NN', 'IN', '$', 'CD', 'CD'])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_markables[1][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the markables now have the `tags` member populated with the part-of-speech tags for each token in the `string` field.\n",
    "\n",
    "**Deliverable 2.6** Now implement a new matcher, `coref_rules.match_on_content()`. Your code should match $m_a$ and $m_i$ iff all content words are identical. It should also enforce the \"no overlap\" restriction defined above.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_match_content_f1_d2_6()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5565\tR: 0.4374\tP:0.7647\n"
     ]
    }
   ],
   "source": [
    "content_matcher = coref_rules.make_resolver(coref_rules.match_on_content)\n",
    "coref.eval_on_dataset(content_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5872\tCEAF: 0.4725\tMUC: 0.5727\tAverage: 0.5441\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(content_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally getting some headway on those metrics!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 2.7** Run the code blocks below to output predictions for the dev and test data.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_dev_acc_f1_d2_7(), test_test_acc_f1_d2_7()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# run once, then remove/comment out cell after dir is created\n",
    "!mkdir predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "coref.write_predictions(coref_rules.make_resolver(coref_rules.match_on_content),\n",
    "                        all_markables_dev,\n",
    "                        'predictions/rules-dev.preds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5463\tR: 0.3960\tP:0.8806\n"
     ]
    }
   ],
   "source": [
    "f,r,p = coref.eval_predictions('predictions/rules-dev.preds',all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5768\tCEAF: 0.4045\tMUC: 0.5463\tAverage: 0.5092\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(content_matcher, all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "coref.write_predictions(coref_rules.make_resolver(coref_rules.match_on_content),\n",
    "                        all_markables_te,\n",
    "                        'predictions/rules-test.preds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.4991\tR: 0.3811\tP:0.7231\n",
      "B-Cubed: 0.5522\tCEAF: 0.4650\tMUC: 0.5168\tAverage: 0.5113\n"
     ]
    }
   ],
   "source": [
    "# students can't run this (well, they can but it'll match against a full-cluster match)\n",
    "coref.eval_predictions('predictions/rules-test.preds', all_markables_te);\n",
    "coref_metrics(content_matcher, all_markables_te);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3: Machine learning for coreference resolution\n",
    "\n",
    "You will now implement coreference resolution using the mention-ranking model. Let's start by implementing some features.\n",
    "\n",
    "**Deliverable 3.1** Implement `coref_features.minimal_features`, using the rules you wrote from `coref_rules.` This should be a function that takes a list of markables, and indices for two mentions, and returns a dict with features and counts. Include the following features:\n",
    "\n",
    "- `exact-match`\n",
    "- `last-token-match`\n",
    "- `content-match`\n",
    "- `crossover`: value of 1 iff the mentions overlap\n",
    "- `new-entity`: value of 1 iff i=j\n",
    "\n",
    "For the first four features, you should call your code from coref_rules directly.\n",
    "(0.25 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.test_minimal_features_d3_1()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 Markable(string=['South', 'Korea'], entity='set_971', start_token=0, end_token=2, tags=['NNP', 'NNP'])\n",
      "1 Markable(string=['a', 'trade', 'deficit', 'of', '$', '101', 'million'], entity='set_972', start_token=3, end_token=10, tags=['DT', 'NN', 'NN', 'IN', '$', 'CD', 'CD'])\n",
      "2 Markable(string=['October'], entity='set_973', start_token=11, end_token=12, tags=['NNP'])\n",
      "3 Markable(string=['the', 'country', \"'s\", 'economic', 'sluggishness'], entity='set_974', start_token=14, end_token=19, tags=['DT', 'NN', 'POS', 'JJ', 'NN'])\n",
      "4 Markable(string=['country'], entity='set_971', start_token=15, end_token=16, tags=['NN'])\n",
      "5 Markable(string=['government', 'figures'], entity='set_975', start_token=22, end_token=24, tags=['NN', 'NNS'])\n",
      "6 Markable(string=['Wednesday'], entity='set_976', start_token=25, end_token=26, tags=['NNP'])\n",
      "7 Markable(string=['Preliminary', 'tallies'], entity='set_977', start_token=27, end_token=29, tags=['JJ', 'NNS'])\n",
      "8 Markable(string=['the', 'Trade', 'and', 'Industry', 'Ministry'], entity='set_978', start_token=30, end_token=35, tags=['DT', 'NNP', 'CC', 'NNP', 'NNP'])\n",
      "9 Markable(string=['another', 'trade', 'deficit', 'the', 'fifth', 'monthly', 'setback', 'this', 'year'], entity='set_979', start_token=36, end_token=48, tags=['DT', 'NN', 'NN', 'IN', 'NNP', ',', 'DT', 'JJ', 'JJ', 'NN', 'DT', 'NN'])\n",
      "10 Markable(string=['October'], entity='set_973', start_token=40, end_token=41, tags=['NNP'])\n",
      "11 Markable(string=['this', 'year'], entity='set_980', start_token=46, end_token=48, tags=['DT', 'NN'])\n",
      "12 Markable(string=['a', 'cloud'], entity='set_981', start_token=50, end_token=52, tags=['DT', 'NN'])\n",
      "13 Markable(string=['South', 'Korea'], entity='set_971', start_token=53, end_token=55, tags=['NNP', 'NNP'])\n",
      "14 Markable(string=['South', 'Korea', \"'s\", 'export-oriented', 'economy'], entity='set_982', start_token=53, end_token=58, tags=['NNP', 'NNP', 'POS', 'JJ', 'NN'])\n"
     ]
    }
   ],
   "source": [
    "min_features = ['exact-match', 'last-token-match', 'content-match', 'crossover', 'new-entity']\n",
    "for i, markable in enumerate(all_markables[1][:15]):\n",
    "    print(i, markable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "defaultdict(<class 'float'>, {})\n",
      "defaultdict(<class 'float'>, {'exact-match': 1.0, 'last-token-match': 1.0, 'content-match': 1.0})\n",
      "defaultdict(<class 'float'>, {'crossover': 1.0})\n",
      "defaultdict(<class 'float'>, {'new-entity': 1.0})\n",
      "defaultdict(<class 'float'>, {'exact-match': 1.0, 'last-token-match': 1.0, 'content-match': 1.0})\n"
     ]
    }
   ],
   "source": [
    "print(coref_features.minimal_features(all_markables[1],0,1))\n",
    "print(coref_features.minimal_features(all_markables[1],0,13))\n",
    "print(coref_features.minimal_features(all_markables[1],13,14))\n",
    "print(coref_features.minimal_features(all_markables[1],6,6))\n",
    "print(coref_features.minimal_features(all_markables[1],2,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 3.2** You will now use these features in a simple feedforward neural net. Using pytorch, implement the `__init__()` and `forward()` functions in `coref_learning.FFCoref` which will be composed of two linear layers separated by a tanh nonlinearity, producing a score for each possible antecedent.\n",
    "\n",
    "Later we will use this scoring function to select the most probable antecendent.\n",
    "(1 point)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_ffcoref_d3_2()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "COREF_FF_HIDDEN = 5 # dimension for hidden layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      "-0.1148\n",
      "[torch.FloatTensor of size 1]\n",
      "\n",
      "Variable containing:\n",
      "-0.3763\n",
      "[torch.FloatTensor of size 1]\n",
      "\n",
      "Variable containing:\n",
      "-0.3763\n",
      "[torch.FloatTensor of size 1]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984) # DO NOT CHANGE\n",
    "coref_ff = coref_learning.FFCoref(min_features, COREF_FF_HIDDEN)\n",
    "\n",
    "# scores for single mention pairs, no backprop\n",
    "print(coref_ff(coref_features.minimal_features(all_markables[1],0,1)))\n",
    "print(coref_ff(coref_features.minimal_features(all_markables[1],0,13)))\n",
    "print(coref_ff(coref_features.minimal_features(all_markables[1],2,10)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 3.3** Given a markable and all its possible antecedents, we now wish to feed all the scores into a softmax layer and attempt the highest possible sum of likelihoods for antecedents representing the correct entity. Implement `FFCoref.score_instance()` in order to do so.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_ffcoref_score_instance_d3_3()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      "\n",
      "Columns 0 to 9 \n",
      "-0.3763 -0.1148 -0.1148 -0.1148 -0.1148 -0.1148 -0.1148 -0.1148 -0.1148 -0.1148\n",
      "\n",
      "Columns 10 to 13 \n",
      "-0.1148 -0.1148 -0.1148 -0.0424\n",
      "[torch.FloatTensor of size 1x14]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "i_scores = coref_ff.score_instance(all_markables[1], coref_features.minimal_features, 13)\n",
    "print(i_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In inference time, all we need is to use the above function and report the antecedent for each markable. In training time, we will use an objective based on the **hinge margin-loss** function. Our variant will require the highest-scoring false candidate (according to the true entity annotations) to score lower than the highest-scoring true candidate *by a margin*:\n",
    "\n",
    "$$L_{m_i} = \\{ \\text{max}_{a:m_a\\notin A(m_i)} s(m_i,m_a) + M - \\text{max}_{a:m_a\\in A(m_i)} s(m_i,m_a) \\}_{+}$$\n",
    "\n",
    "Where $s$ is the score from the previous deliverable, $A(m)$ denotes the set of true antecedents for $m$, $M$ is our margin, and the $+$ subscript indicates that negative values are replaced by $0$ (since this is a hinge loss).\n",
    "\n",
    "**Deliverable 3.4**\n",
    "Implement the helper function `FFCoref.instance_top_scores()` which supplies the arguments for this loss function.\n",
    "**Note** the special cases where:\n",
    "- Only true candidates exist (we're in the first cluster) - the trainer will have to skip these. Return `None`s.\n",
    "- Only false candidates exist - this actually means we're in a new cluster.\n",
    "\n",
    "(0.5 points)\n",
    "\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_ffcoref_score_instance_top_scores_d3_4()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      "-0.1148\n",
      "-0.0424\n",
      "[torch.FloatTensor of size 2]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "best_true_score, best_false_score = coref_ff.instance_top_scores(all_markables[1], coref_features.minimal_features, 13, 4)\n",
    "print(torch.cat([best_true_score, best_false_score], 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks like we're ready to train our classifier!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss = 0.9191775133819172\n",
      "Loss = 0.9106459462424766\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984) # DO NOT CHANGE\n",
    "coref_ff = coref_learning.FFCoref(min_features, COREF_FF_HIDDEN)\n",
    "optimizer = optim.SGD(coref_ff.parameters(), lr=ETA_0)\n",
    "coref_learning.train(coref_ff, optimizer, all_markables, coref_features.minimal_features, margin=1.0, epochs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5310\tR: 0.4176\tP:0.7287\n"
     ]
    }
   ],
   "source": [
    "# training set results\n",
    "coref_learning.evaluate(coref_ff, all_markables, coref_features.minimal_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5505\tR: 0.4027\tP:0.8696\n"
     ]
    }
   ],
   "source": [
    "# dev set\n",
    "coref_learning.evaluate(coref_ff, all_markables_dev, coref_features.minimal_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5869\tCEAF: 0.4718\tMUC: 0.5723\tAverage: 0.5436\n"
     ]
    }
   ],
   "source": [
    "# standard metrics on training set\n",
    "ff_matcher = coref_learning.make_resolver(coref_features.minimal_features, coref_ff)\n",
    "coref_metrics(ff_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5522\tCEAF: 0.4650\tMUC: 0.5168\tAverage: 0.5113\n"
     ]
    }
   ],
   "source": [
    "#test set, students can't run\n",
    "coref_metrics(ff_matcher, all_markables_te);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 3.5** We can add more features to try and better capture relations between markables.\n",
    "\n",
    "Implement distance features in `coref_features.distance_features()`, measuring the mention distance and the token distance. Specifically:\n",
    "\n",
    "- **Mention distance** is number of intervening mentions between i and j, $i-j$.\n",
    "- **Token distance** is number of tokens between the start of i and the end of j.\n",
    "\n",
    "These should be binary features, up to a maximum distance of 10 for tokens / 5 for mentions, with the final feature indicating distance of 10/5 and above, respectively. The desired behavior is shown below.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_distance_features_d3_5()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 Markable(string=['South', 'Korea'], entity='set_971', start_token=0, end_token=2, tags=['NNP', 'NNP'])\n",
      "1 Markable(string=['a', 'trade', 'deficit', 'of', '$', '101', 'million'], entity='set_972', start_token=3, end_token=10, tags=['DT', 'NN', 'NN', 'IN', '$', 'CD', 'CD'])\n",
      "2 Markable(string=['October'], entity='set_973', start_token=11, end_token=12, tags=['NNP'])\n",
      "3 Markable(string=['the', 'country', \"'s\", 'economic', 'sluggishness'], entity='set_974', start_token=14, end_token=19, tags=['DT', 'NN', 'POS', 'JJ', 'NN'])\n"
     ]
    }
   ],
   "source": [
    "for i, markable_i in enumerate(all_markables[1][:4]):\n",
    "    print(i, markable_i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "defaultdict(<class 'float'>, {})\n",
      "defaultdict(<class 'float'>, {'mention-distance-1': 1, 'token-distance-1': 1})\n",
      "defaultdict(<class 'float'>, {'mention-distance-2': 1, 'token-distance-9': 1})\n",
      "defaultdict(<class 'float'>, {'mention-distance-2': 1, 'token-distance-4': 1})\n",
      "defaultdict(<class 'float'>, {'mention-distance-5': 1, 'token-distance-10': 1})\n"
     ]
    }
   ],
   "source": [
    "print(coref_features.distance_features(all_markables[1],0,0))\n",
    "print(coref_features.distance_features(all_markables[1],0,1))\n",
    "print(coref_features.distance_features(all_markables[1],0,2))\n",
    "print(coref_features.distance_features(all_markables[1],1,3))\n",
    "print(coref_features.distance_features(all_markables[1],0,30))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 3.6** Implement `coref_features.make_feature_union()`, which should take a list of feature functions, and return a function that computes the union of all features in the list. You can assume the feature functions don't use the same name for any feature.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_coref.py:test_feature_union_d3_6()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "joint_feats = coref_features.make_feature_union([coref_features.minimal_features,\n",
    "                                                 coref_features.distance_features])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "defaultdict(<class 'float'>, {'mention-distance-2': 1, 'token-distance-4': 1})\n",
      "defaultdict(<class 'float'>, {'mention-distance-3': 1, 'token-distance-10': 1})\n",
      "defaultdict(<class 'float'>, {'mention-distance-5': 1, 'token-distance-10': 1})\n",
      "defaultdict(<class 'float'>, {'new-entity': 1.0})\n"
     ]
    }
   ],
   "source": [
    "print(joint_feats(all_markables[1],1,3))\n",
    "print(joint_feats(all_markables[1],0,3))\n",
    "print(joint_feats(all_markables[1],0,7))\n",
    "print(joint_feats(all_markables[1],10,10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "min_features_and_distances = min_features\\\n",
    "                                   + [f'mention-distance-{i}' for i in range(1,6)]\\\n",
    "                                   + [f'token-distance-{i}' for i in range(1,11)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss = 0.9255589622700262\n",
      "Loss = 0.8616410730355986\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984) # DO NOT CHANGE\n",
    "coref_ff_w_distances = coref_learning.FFCoref(min_features_and_distances, 50) # we need more hidden units now\n",
    "optimizer = optim.SGD(coref_ff_w_distances.parameters(), lr=ETA_0)\n",
    "coref_learning.train(coref_ff_w_distances, optimizer, all_markables, joint_feats, epochs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5424\tR: 0.4118\tP:0.7942\n"
     ]
    }
   ],
   "source": [
    "coref_learning.evaluate(coref_ff_w_distances, all_markables, joint_feats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5780\tCEAF: 0.4572\tMUC: 0.5623\tAverage: 0.5325\n"
     ]
    }
   ],
   "source": [
    "ff_w_dist_matcher = coref_learning.make_resolver(joint_feats, coref_ff_w_distances)\n",
    "coref_metrics(ff_w_dist_matcher, all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that our basic F metric got slightly better, while the average standard metric got slightly worse due to reduced MUC and CEAF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5213\tR: 0.3691\tP:0.8871\n"
     ]
    }
   ],
   "source": [
    "coref.write_predictions(ff_w_dist_matcher,\n",
    "                        all_markables_dev,\n",
    "                        'predictions/ff-dev.preds')\n",
    "coref.eval_predictions('predictions/ff-dev.preds', all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5641\tCEAF: 0.3940\tMUC: 0.5308\tAverage: 0.4963\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(ff_w_dist_matcher, all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "coref.write_predictions(ff_w_dist_matcher,\n",
    "                        all_markables_te,\n",
    "                        'predictions/ff-test.preds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.4799\tR: 0.3541\tP:0.7443\n",
      "B-Cubed: 0.5414\tCEAF: 0.4356\tMUC: 0.4908\tAverage: 0.4893\n"
     ]
    }
   ],
   "source": [
    "# students can't run this\n",
    "coref.eval_predictions('predictions/ff-test.preds', all_markables_te);\n",
    "coref_metrics(ff_w_dist_matcher, all_markables_te);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4: Sequential Text Represenation\n",
    "\n",
    "In this section, we will find out whether neural representations of our text can help find coreferents.\n",
    "\n",
    "The main idea is to run a bidirectional LSTM model, which you already have implemented from previous problem sets, and use the resulting hidden states to form representations of the markables. These will be fed into a feedforward classifier similar to the one from the previous section, except that the match features will also be embedded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3530\n"
     ]
    }
   ],
   "source": [
    "# Preparing the vocabulary for a word-to-index dictionary necessary for the initial embeddings table\n",
    "vocab = set()\n",
    "for doc in all_words + all_words_dev + all_words_test:\n",
    "    vocab.update(doc)\n",
    "vocab = sorted(list(vocab))\n",
    "word_to_ix = {w:i for i,w in enumerate(vocab)}\n",
    "print(len(vocab))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 4.1**\n",
    "The first part will be very similar to code from PS3, which you may reuse. Implement `neural_net.BiLSTMWordEmbedding` as a word  embedding lookup table followed by a bi-directional LSTM which runs on a text (here, the entire document) and outputs the hidden state from the LSTM as a contextual embedding for each word in it.\n",
    "(1 point)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_bilstm_embedding_d4_1()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "WORD_EMB_DIM = 64\n",
    "WORD_LSTM_EMB_DIM = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "McDermott International Inc. said its Babcock & Wilcox unit completed the sale of its Bailey Controls Operations ... \n",
      "\n",
      "its \n",
      " Variable containing:\n",
      "-0.1957\n",
      "-0.2839\n",
      " 0.2615\n",
      "-0.0888\n",
      "-0.1205\n",
      "[torch.FloatTensor of size 5]\n",
      "\n",
      "its \n",
      " Variable containing:\n",
      "-0.1274\n",
      "-0.1916\n",
      " 0.2225\n",
      "-0.1166\n",
      "-0.0630\n",
      "[torch.FloatTensor of size 5]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984)\n",
    "word_lstm = neural_net.BiLSTMWordEmbedding(word_to_ix, WORD_EMB_DIM, WORD_LSTM_EMB_DIM, 1, 0.5)\n",
    "embs = word_lstm(all_words[0])\n",
    "print(' '.join(all_words[0][:17] + ['...']), '\\n')\n",
    "print(all_words[0][4], '\\n', embs[4][0][:5])\n",
    "print(all_words[0][13], '\\n', embs[13][0][:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see how the same word type (*its*) is assigned different embeddings based on its context in the document."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention Model\n",
    "\n",
    "Our markable embeddings will be trained using an **Attention Model** which accepts the embeddings for the words $\\{w_i\\}$ in a markable and outputs a single vector that \"attends\" to the single vectors according to what it believes is their importance.\n",
    "This concept, [originally used](https://arxiv.org/abs/1409.0473) for sequence-to-sequence models such as Machine Translation, is applied for our task as yet another attempt to find the crucial part of a markable, like we did for the head-finding heuristic and for the content-word matching. While those were \"hard\" techniques, yes-or-no for each token, here we're applying a \"soft\" weighting that still assigns all the words in the text some significance.\n",
    "\n",
    "Practically, our model will train a vector parameter of the same embedding size as the BiLSTM output, $\\vec{u}$. Each word's embedding in the input span $\\vec{e_i}$ will be multiplied (dot-product) with $\\vec{u}$ to assign it a scalar weight, $a_i$. Finally each embedding will be multiplied by its normalized (softmaxed) weight $\\alpha_i$, and the sum of these weighted vectors will be our markable's output embedding, $\\vec{e_m}$:\n",
    "\n",
    "![Attention equations](att.PNG \"Attention equations\")\n",
    "\n",
    "**Deliverable 4.2**\n",
    "Implement `neural_net.AttentionBasedMarkableEmbedding`.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_embedding_attention_d4_2()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['McDermott', 'International', 'Inc.'] set_356 \n",
      " Variable containing:\n",
      "-0.1441\n",
      "-0.1028\n",
      "-0.1215\n",
      "[torch.FloatTensor of size 3]\n",
      "\n",
      "['its'] set_356 \n",
      " Variable containing:\n",
      "-0.1957\n",
      "-0.2839\n",
      " 0.2615\n",
      "[torch.FloatTensor of size 3]\n",
      "\n",
      "['its'] set_357 \n",
      " Variable containing:\n",
      "-0.1274\n",
      "-0.1916\n",
      " 0.2225\n",
      "[torch.FloatTensor of size 3]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984)\n",
    "attn_layer = neural_net.AttentionBasedMarkableEmbedding(WORD_LSTM_EMB_DIM)\n",
    "mark_embs = [attn_layer(embs, m) for m in all_markables[0]]\n",
    "for j in [0, 1, 4]:\n",
    "    print(all_markables[0][j].string, all_markables[0][j].entity, '\\n', mark_embs[j][:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will score a markable pair based on the following features:\n",
    "\n",
    "1. Each markable's attended embedding\n",
    "1. A low-dimension embedding for each of the pairwise features we've extracted in the previous sections. Since they are boolean, each will have an embedding for its ''false'' state and one for its ''true'' state.\n",
    "\n",
    "First, let's implement a quick extractor for positive-valued features from a mention pair."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_positive_feats(doc, i, a, feats=coref_features.minimal_features):\n",
    "    return [k for k,v in feats(doc,i,a).items() if v > 0.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['exact-match', 'last-token-match', 'content-match']"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_positive_feats(all_markables[1], 0, 13)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will concatenate all of these embeddings together and use them as input in a two-layer feedforward network (with ReLU nonlinearity) which will produce a scalar score for markable match.\n",
    "\n",
    "**Deliverable 4.3**\n",
    "Implement `__init__()` and `forward()` in `neural_net.SequentialScorer`.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_sequential_scorer_d4_3()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       "1.00000e-02 *\n",
       " -6.5129\n",
       "[torch.FloatTensor of size 1x1]"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BOOLEAN_FEATURE_DIM = 6\n",
    "SCORER_HIDDEN_DIM = 164\n",
    "\n",
    "# starting with just one document\n",
    "torch.manual_seed(1984)\n",
    "word_lstm1 = neural_net.BiLSTMWordEmbedding(word_to_ix, WORD_EMB_DIM, WORD_LSTM_EMB_DIM, 1, 0.5)\n",
    "embs1 = word_lstm1(all_words[1])\n",
    "attn_layer1 = neural_net.AttentionBasedMarkableEmbedding(WORD_LSTM_EMB_DIM)\n",
    "mkbls1 = all_markables[1]\n",
    "mark_embs1 = [attn_layer(embs1, m) for m in mkbls1]\n",
    "scorer = neural_net.SequentialScorer(WORD_LSTM_EMB_DIM, min_features, BOOLEAN_FEATURE_DIM, SCORER_HIDDEN_DIM)\n",
    "scorer(mark_embs1[13], mark_embs1[0], get_positive_feats(mkbls1, 13, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 4.4**\n",
    "Implement `score_instance()` and `instance_top_scores()` in `neural_net.SequentialScorer`. Their purpose is the same as the one in `FFCoref`, but they require the extra embeddings parameter. The former will require some changes to adapt to the different `forward()`, but the latter can be identical to its correlate in `FFCoref` if implemented correctly.\n",
    "(1 point)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.py:test_sequential_scorer_score_instance_d4_4(), test_sequential_scorer_instance_top_scores_d4_4()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       "\n",
       "Columns 0 to 9 \n",
       "1.00000e-02 *\n",
       " -6.5129 -3.6650 -2.7447 -4.9805 -3.6708 -3.9170 -3.0245 -2.2054 -3.0073 -2.2837\n",
       "\n",
       "Columns 10 to 13 \n",
       "1.00000e-02 *\n",
       " -2.4808 -5.2409 -1.9311 -1.6140\n",
       "[torch.FloatTensor of size 1x14]"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scorer.score_instance(mark_embs1, all_markables[1], 13, coref_features.minimal_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to the length of our documents and number of parameters, torch may not work properly on the entire text. We'll truncate the files before we train and use only the minimal pairwise feature space, causing our performance to be suboptimal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 31 complete.\n",
      "Document losses = 0.93932, 0.62971, 1.57278, 0.91849, 0.89681, 0.99915, 0.93578, 0.59647, 0.47273, 0.96949, 0.80176, 0.81332, 0.90159, 0.91200, 0.74807, 0.68139, 0.50189, 0.59706, 0.90399, 0.27876, 1.22716, 1.08047\n",
      "Overall loss = 0.83621\n",
      "Epoch 31 complete.\n",
      "Document losses = 0.95553, 0.89429, 0.66035, 0.46930, 0.78851, 0.93100, 0.84361, 0.83657, 0.39137, 0.71947, 0.78565, 1.10904, 0.57058, 1.02704, 0.81013, 0.43853, 0.40327, 0.76322, 0.62335, 1.40899, 0.73699, 1.09462\n",
      "Overall loss = 0.77141\n",
      "Epoch 31 complete.\n",
      "Document losses = 0.69691, 0.31914, 1.00857, 0.64840, 0.77790, 0.90615, 0.81879, 0.94573, 0.48738, 0.63369, 0.66208, 0.62734, 0.56925, 0.57159, 0.65879, 0.43431, 0.65710, 0.76916, 0.47507, 0.22506, 0.84899, 1.04005\n",
      "Overall loss = 0.66926\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984)\n",
    "tr_word_lstm = neural_net.BiLSTMWordEmbedding(word_to_ix, WORD_EMB_DIM, WORD_LSTM_EMB_DIM, 1, 0.2)\n",
    "tr_attn_layer = neural_net.AttentionBasedMarkableEmbedding(WORD_LSTM_EMB_DIM)\n",
    "tr_scorer = neural_net.SequentialScorer(WORD_LSTM_EMB_DIM, min_features, BOOLEAN_FEATURE_DIM, SCORER_HIDDEN_DIM)\n",
    "optimizer = optim.SGD(list(tr_word_lstm.parameters()) + list(tr_attn_layer.parameters()) + list(tr_scorer.parameters()), lr=ETA_0)\n",
    "neural_net.train(tr_word_lstm, tr_attn_layer, tr_scorer,\\\n",
    "                 optimizer, all_words, all_markables, coref_features.minimal_features, word_limit=150, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's evaluate on the entire dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5207\tR: 0.3944\tP:0.7658\n"
     ]
    }
   ],
   "source": [
    "tr_resolver = neural_net.evaluate(tr_word_lstm, tr_attn_layer, tr_scorer, all_words, all_markables, coref_features.minimal_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5756\tCEAF: 0.4551\tMUC: 0.5605\tAverage: 0.5304\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(tr_resolver, all_markables);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5403\tR: 0.3826\tP:0.9194\n"
     ]
    }
   ],
   "source": [
    "dv_resolver = neural_net.evaluate(tr_word_lstm, tr_attn_layer, tr_scorer, all_words_dev, all_markables_dev, coref_features.minimal_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5653\tCEAF: 0.3966\tMUC: 0.5403\tAverage: 0.5007\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(dv_resolver, all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.4630\tR: 0.3378\tP:0.7353\n",
      "B-Cubed: 0.5334\tCEAF: 0.4225\tMUC: 0.4778\tAverage: 0.4779\n"
     ]
    }
   ],
   "source": [
    "# students can't run this\n",
    "te_resolver = neural_net.evaluate(tr_word_lstm, tr_attn_layer, tr_scorer, all_words_test, all_markables_te, coref_features.minimal_features)\n",
    "coref_metrics(te_resolver, all_markables_te);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pretrained Word Embeddings\n",
    "\n",
    "**Deliverable 4.5**\n",
    "Implement `utils.initialize_with_pretrained()`. Start by copying from your implementation in PS3.\n",
    "(0.5 points)\n",
    "\n",
    "* **Test:** `tests\\test_neural_coref.test_pretrain_embeddings_d4_5()`\n",
    "\n",
    "**Note** that there is a new pretrained file in the `data` folder. Although from the same original source, it is trimmed to the vocabulary in our new dataset, so don't use the same file from PS3. In addition to this attribute, it includes a special token called **&lt;UNK&gt;**. You should use its assigned vector to initialize vectors for all unknown words in the dataset.\n",
    "\n",
    "Later, if you're interested in improving your model's performance, you may want to know that there are more special tokens with trained vectors in the data file:\n",
    "* **&lt;S&gt;** signifies the beginning of a sentence.\n",
    "* **&lt;/S&gt;** signifies the end of a sentence.\n",
    "* **&lt;PAD&gt;** is used to pad short sentences (this one probably won't be useful)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pret_embs = pickle.load(open('data/pretrained-embeds-coref.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.34177107 -0.1528549  -0.06971747  0.12536518  0.31670848]\n"
     ]
    }
   ],
   "source": [
    "print(pret_embs['Fujitsu'][:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.2400684   0.0170402  -0.5328812   0.16161029 -0.03450763]\n"
     ]
    }
   ],
   "source": [
    "print(pret_embs['<UNK>'][:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 31 complete.\n",
      "Document losses = 0.93080, 0.65143, 1.54045, 0.94181, 0.92525, 0.99127, 0.95597, 0.78670, 0.58961, 1.09405, 0.64312, 0.75519, 0.83328, 0.82583, 0.49011, 0.64926, 0.60697, 0.60192, 0.76932, 0.31302, 1.40232, 1.03551\n",
      "Overall loss = 0.83530\n",
      "Epoch 31 complete.\n",
      "Document losses = 0.97924, 0.95806, 0.87604, 0.71278, 1.11671, 1.07097, 0.93572, 0.83217, 0.40688, 0.77817, 0.62955, 0.56910, 0.63388, 0.80036, 0.45497, 0.65841, 0.79375, 0.69298, 0.77241, 0.18915, 1.58285, 1.09280\n",
      "Overall loss = 0.80113\n",
      "Epoch 31 complete.\n",
      "Document losses = 1.00449, 0.87274, 0.61690, 0.45452, 0.78579, 0.86951, 0.79216, 0.57742, 0.55570, 0.53367, 0.76431, 0.51676, 0.55175, 0.94359, 0.84593, 0.43401, 0.36180, 0.78888, 0.64513, 1.27888, 0.68404, 1.03125\n",
      "Overall loss = 0.71666\n",
      "Epoch 31 complete.\n",
      "Document losses = 0.85349, 0.21191, 1.25798, 0.74668, 1.24524, 1.03969, 0.88672, 0.67221, 0.36734, 0.55026, 0.64822, 0.66793, 0.58037, 0.60304, 0.71610, 0.44662, 0.64840, 0.82709, 0.51960, 0.18410, 0.88014, 1.05845\n",
      "Overall loss = 0.70438\n",
      "Epoch 31 complete.\n",
      "Document losses = 0.83462, 0.19412, 0.93821, 0.63640, 0.76690, 0.98298, 0.91020, 0.77722, 0.32671, 0.66741, 0.82634, 0.75198, 0.48915, 0.59004, 0.64925, 0.43700, 0.61274, 0.77087, 0.48063, 0.17491, 0.84386, 1.29280\n",
      "Overall loss = 0.66981\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984)\n",
    "tr_pt_word_lstm = neural_net.BiLSTMWordEmbedding(word_to_ix, WORD_EMB_DIM, WORD_LSTM_EMB_DIM, 1, 0.2)\n",
    "utils.initialize_with_pretrained(pret_embs, tr_pt_word_lstm)\n",
    "tr_pt_attn_layer = neural_net.AttentionBasedMarkableEmbedding(WORD_LSTM_EMB_DIM)\n",
    "tr_pt_scorer = neural_net.SequentialScorer(WORD_LSTM_EMB_DIM, min_features, BOOLEAN_FEATURE_DIM, SCORER_HIDDEN_DIM)\n",
    "optimizer = optim.SGD(list(tr_pt_word_lstm.parameters()) + list(tr_pt_attn_layer.parameters()) + list(tr_pt_scorer.parameters()), lr=ETA_0)\n",
    "neural_net.train(tr_pt_word_lstm, tr_pt_attn_layer, tr_pt_scorer,\\\n",
    "                 optimizer, all_words, all_markables, coref_features.minimal_features, word_limit=150, epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5253\tR: 0.3979\tP:0.7725\n"
     ]
    }
   ],
   "source": [
    "tr_pt_resolver = neural_net.evaluate(tr_pt_word_lstm, tr_pt_attn_layer, tr_pt_scorer, all_words, all_markables, coref_features.minimal_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5756\tCEAF: 0.4551\tMUC: 0.5605\tAverage: 0.5304\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(tr_pt_resolver, all_markables);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5308\tR: 0.3758\tP:0.9032\n",
      "B-Cubed: 0.5653\tCEAF: 0.3966\tMUC: 0.5403\tAverage: 0.5007\n"
     ]
    }
   ],
   "source": [
    "tr_pt_resolver_dev = neural_net.evaluate(tr_pt_word_lstm, tr_pt_attn_layer, tr_pt_scorer, all_words_dev, all_markables_dev, coref_features.minimal_features)\n",
    "coref.write_predictions(tr_pt_resolver_dev,\n",
    "                        all_markables_dev,\n",
    "                        'predictions/nn-dev.preds');\n",
    "coref_metrics(tr_pt_resolver_dev, all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.4556\tR: 0.3324\tP:0.7235\n",
      "B-Cubed: 0.5334\tCEAF: 0.4225\tMUC: 0.4778\tAverage: 0.4779\n"
     ]
    }
   ],
   "source": [
    "# students can't run this\n",
    "tr_pt_resolver_te = neural_net.evaluate(tr_pt_word_lstm, tr_pt_attn_layer, tr_pt_scorer, all_words_test, all_markables_te, coref_features.minimal_features)\n",
    "coref.write_predictions(tr_pt_resolver_te,\n",
    "                        all_markables_te,\n",
    "                        'predictions/nn-test.preds');\n",
    "coref_metrics(tr_pt_resolver_te, all_markables_te);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Deliverable 4.6** (7650 only; 4650 optional)\n",
    "\n",
    "To match nominals, it is often necessary to capture **semantics**. Find a paper (in ACL, NAACL, EACL, EMNLP, or TACL, since 2007) that attempts to use semantic analysis to do nominal coreference, and explain:\n",
    "\n",
    "- What form of semantics they are trying to capture (e.g., synonymy, hypernymy, predicate-argument, distributional)\n",
    "- How they formalize semantics into features, constraints, or some other preference\n",
    "- How much it helps\n",
    "\n",
    "Put your answer in `text-answers.md`.\n",
    "(1 point)\n",
    "\n",
    "As usual, if you are in 4650 and you do this problem, you will be graded on the 7650 rubric."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Domain Adaptation (no deliverables)\n",
    "\n",
    "Our dataset has a second part, which is a corpus of fairy tales, rather than news stories.\n",
    "\n",
    "Let's take advantage of this setup to see how well our WSJ-trained model does over the fairy tale data, as opposed to a model trained on the fairy tales themselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ft_dv_dir = os.path.join('data','tales','dev')\n",
    "ft_tr_dir = os.path.join('data','tales','train')\n",
    "ft_te_dir = os.path.join('data','tales','test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ft_all_markables, ft_all_words = coref.read_dataset(ft_tr_dir, tagger=pos_tag)\n",
    "ft_all_markables_dev, ft_all_words_dev = coref.read_dataset(ft_dv_dir, tagger=pos_tag)\n",
    "ft_all_markables_te, ft_all_words_te = coref.read_dataset(ft_te_dir, tagger=pos_tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.6511\tR: 0.5545\tP:0.7884\n"
     ]
    }
   ],
   "source": [
    "coref.eval_on_dataset(exact_matcher, ft_all_markables);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The exact matcher is getting much higher numbers on this dataset than on WSJ. Let's train an ML system and see the differences. We'll use `FFCoref` from section 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss = 0.9428621476767014\n",
      "Loss = 0.9798036302839007\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1984) # DO NOT CHANGE\n",
    "coref_ff_fairy = coref_learning.FFCoref(min_features_and_distances, 50)\n",
    "optimizer = optim.SGD(coref_ff_fairy.parameters(), lr=ETA_0)\n",
    "coref_learning.train(coref_ff_fairy, optimizer, ft_all_markables, joint_feats, epochs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F: 0.5583\tR: 0.4747\tP:0.6776\n"
     ]
    }
   ],
   "source": [
    "# in-domain trained, tested on fairy tale data\n",
    "coref_learning.evaluate(coref_ff_fairy, ft_all_markables_dev, joint_feats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.4673\tCEAF: 0.3698\tMUC: 0.6070\tAverage: 0.4814\n"
     ]
    }
   ],
   "source": [
    "ft_ff_matcher = coref_learning.make_resolver(joint_feats, coref_ff_fairy)\n",
    "coref_metrics(ft_ff_matcher, ft_all_markables_dev);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the **in-domain** dev set scores for the fairy tale dataset.\n",
    "\n",
    "Recall that the in-domain numbers for the WSJ portion were the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5641\tCEAF: 0.3940\tMUC: 0.5308\tAverage: 0.4963\n"
     ]
    }
   ],
   "source": [
    "coref_metrics(ff_w_dist_matcher, all_markables_dev);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can see how the **cross-domain** results look. Can a model trained on news data be reliably used in a fairy-tale setting? How about the converse?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.5329\tCEAF: 0.4101\tMUC: 0.4692\tAverage: 0.4707\n"
     ]
    }
   ],
   "source": [
    "# trained on fairy, tested on WSJ\n",
    "coref_metrics(ft_ff_matcher, all_markables_dev);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B-Cubed: 0.4615\tCEAF: 0.3590\tMUC: 0.6122\tAverage: 0.4776\n"
     ]
    }
   ],
   "source": [
    "# trained on WSJ, tested on fairy\n",
    "coref_metrics(ff_w_dist_matcher, ft_all_markables_dev);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are lower than in-domain, but surprisingly enough not too low. In other tasks this problem is more acute.\n",
    "\n",
    "If you're interested, you may try and join the two training sets together and see where that gets you. You can also use this extra data (only training sets!) for your bakeoff."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Final bakeoff!\n",
    "\n",
    "This is for extra credit only. Ideas for improvements:\n",
    "\n",
    "- Cost-sensitive training to balance precision and recall\n",
    "- Syntax (you can parse all the markables as a preprocessing step)\n",
    "  - Tree distance\n",
    "  - Syntactic parallelism\n",
    "  - Better head matching\n",
    "- Add layers\n",
    "- Better learning\n",
    "- Ensemble (average) multiple models together\n",
    "\n",
    "Feel free to search the research literature (via Google scholar) to get ideas. If you use an idea from another paper, mention the paper (authors, title, and URL) in your comments in `coref_features.py`.\n",
    "\n",
    "As usual, sometimes improvement can also come from tweaking parameters in the neural nets.\n",
    "\n",
    "In section 4, recall that we truncated the training documents - maybe more data (or different portions of each document) can help.\n",
    "\n",
    "To use cuda, pass in use_cuda=True into `neural_net.train()` and `utils.initialize_with_pretrained()`.\n",
    "\n",
    "**Deliverable 6**. Copy the applicable cells from section 3 or 4 by your choice to output predictions for both the dev and test sets.\n",
    "\n",
    "Write your predictions to `predictions/bakeoff-dev.preds` and `predictions/bakeoff-test.preds`.\n",
    "\n",
    "Due to the complex evaluation structure, this will be done from the tarball and not on Kaggle.\n",
    "\n",
    "Scoring:\n",
    "\n",
    "Test set F1 must be greater than .51 to get any extra credit.\n",
    "\n",
    "- Dev F1 > .535: +0.5 points\n",
    "- Dev F1 > .55: +1 points\n",
    "- Dev F1 > .6, Test F1 > .55: +3 points\n",
    "- Best in 4650: +0.5 points\n",
    "- Best in 7650: +0.5 points\n",
    "- Better than best TA/prof system: +0.5 points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
