tutorials/end2end_ned_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# End-to-End NED Tutorial\n",
    "\n",
    "In this tutorial, we walk through how to use Bootleg as an end-to-end pipeline to detect and label entities in a set of sentences. First, we show how to use Bootleg to detect and disambiguate mentions to entities. We then compare to an existing system named TAGME. \n",
    "\n",
    "This tutorial assumes you want to use Bootleg on full datasets. You can also use Bootleg in annotator mode:\n",
    "\n",
    "```\n",
    "pip install bootleg\n",
    "from bootleg.end2end.bootleg_annotator import BootlegAnnotator\n",
    "ann = BootlegAnnotator()\n",
    "ann.label_mentions(\"Bob Dylan release Desire\")[\"titles\"]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To understand how Bootleg performs on more natural language than we find in Wikipedia, we hand label the mentions and corresponding entities in 50 questions sampled from the [Natural Questions dataset (Google)](https://ai.google.com/research/NaturalQuestions). We will evaluate our *uncased* Bootleg model. However, we have manually cased the data in case you want to try our cased model instead.\n",
    "\n",
    "### Requirements\n",
    "\n",
    "You will need to download the following files for this notebook:\n",
    "- Pretrained Bootleg uncased model and config [here](https://bootleg-data.s3-us-west-2.amazonaws.com/models/lateset/bootleg_uncased.tar.gz).\n",
    "- Sample of Natural Questions with hand-labelled entities [here](https://bootleg-data.s3-us-west-2.amazonaws.com/data/lateset/nq.tar.gz)\n",
    "- Entity data [here](https://bootleg-data.s3-us-west-2.amazonaws.com/data/lateset/entity_db.tar.gz)\n",
    "\n",
    "For convenience, you can run the commands below (from the root directory of the repo) to download all the above files and unpack them to `models` and `data` directories. It will take several minutes to download all the files. \n",
    "\n",
    "```\n",
    "    bash tutorials/download_model.sh uncased\n",
    "    bash tutorials/download_data.sh\n",
    "```\n",
    "\n",
    "You can also run directly in this notebook by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2021-10-15 22:22:30--  https://bootleg-data.s3-us-west-2.amazonaws.com/models/latest/bootleg_uncased.tar.gz\n",
      "Resolving bootleg-data.s3-us-west-2.amazonaws.com (bootleg-data.s3-us-west-2.amazonaws.com)... 52.92.128.82\n",
      "Connecting to bootleg-data.s3-us-west-2.amazonaws.com (bootleg-data.s3-us-west-2.amazonaws.com)|52.92.128.82|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 496143916 (473M) [application/x-tar]\n",
      "Saving to: ‘models/bootleg_uncased.tar.gz.1’\n",
      "\n",
      "bootleg_uncased.tar 100%[===================>] 473.16M  25.1MB/s    in 30s     \n",
      "\n",
      "2021-10-15 22:23:06 (15.6 MB/s) - ‘models/bootleg_uncased.tar.gz.1’ saved [496143916/496143916]\n",
      "\n",
      "bootleg_uncased/\n",
      "bootleg_uncased/bootleg_wiki.pth\n",
      "bootleg_uncased/bootleg_config.yaml\n",
      "--2021-10-15 22:23:18--  https://bootleg-data.s3-us-west-2.amazonaws.com/data/latest/nq.tar.gz\n",
      "Resolving bootleg-data.s3-us-west-2.amazonaws.com (bootleg-data.s3-us-west-2.amazonaws.com)... 52.218.252.81\n",
      "Connecting to bootleg-data.s3-us-west-2.amazonaws.com (bootleg-data.s3-us-west-2.amazonaws.com)|52.218.252.81|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 5928 (5.8K) [application/x-tar]\n",
      "Saving to: ‘data/nq.tar.gz.1’\n",
      "\n",
      "nq.tar.gz.1         100%[===================>]   5.79K  --.-KB/s    in 0s      \n",
      "\n",
      "2021-10-15 22:23:18 (196 MB/s) - ‘data/nq.tar.gz.1’ saved [5928/5928]\n",
      "\n",
      "--2021-10-15 22:23:18--  https://bootleg-data.s3-us-west-2.amazonaws.com/data/latest/entity_db.tar.gz\n",
      "Resolving bootleg-data.s3-us-west-2.amazonaws.com (bootleg-data.s3-us-west-2.amazonaws.com)... 52.218.252.81\n",
      "Connecting to bootleg-data.s3-us-west-2.amazonaws.com (bootleg-data.s3-us-west-2.amazonaws.com)|52.218.252.81|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 1862741895 (1.7G) [application/x-tar]\n",
      "Saving to: ‘data/entity_db.tar.gz.1’\n",
      "\n",
      "entity_db.tar.gz.1  100%[===================>]   1.73G  18.0MB/s    in 2m 11s  \n",
      "\n",
      "2021-10-15 22:25:34 (13.5 MB/s) - ‘data/entity_db.tar.gz.1’ saved [1862741895/1862741895]\n",
      "\n",
      "nq/\n",
      "nq/test_50_bootleg.jsonl\n",
      "nq/test_50_tagme.jsonl\n",
      "nq/test_50.jsonl\n",
      "entity_db/\n",
      "entity_db/entity_mappings/\n",
      "entity_db/entity_mappings/alias2qids_unfiltered.json\n",
      "entity_db/entity_mappings/qid2eid.json\n",
      "entity_db/entity_mappings/alias2id_unfiltered.json\n",
      "entity_db/entity_mappings/alias2id.json\n",
      "entity_db/entity_mappings/qid2title.json\n",
      "entity_db/entity_mappings/qid2desc.json\n",
      "entity_db/entity_mappings/alias2qids.json\n",
      "entity_db/entity_mappings/qid2cnt.json\n",
      "entity_db/entity_mappings/config.json\n",
      "entity_db/kg_mappings/\n",
      "entity_db/kg_mappings/qid2relations.json\n",
      "entity_db/kg_mappings/config.json\n",
      "entity_db/kg_mappings/kg_adj.txt\n",
      "entity_db/kg_mappings/relation_vocab.json\n",
      "entity_db/type_mappings/\n",
      "entity_db/type_mappings/relations/\n",
      "entity_db/type_mappings/relations/config.json\n",
      "entity_db/type_mappings/relations/type_vocab.json\n",
      "entity_db/type_mappings/relations/qid2typeids.json\n",
      "entity_db/type_mappings/relations/qid2typenames.json\n",
      "entity_db/type_mappings/hyena/\n",
      "entity_db/type_mappings/hyena/config.json\n",
      "entity_db/type_mappings/hyena/type_vocab.json\n",
      "entity_db/type_mappings/hyena/qid2typeids.json\n",
      "entity_db/type_mappings/hyena/qid2typenames.json\n",
      "entity_db/type_mappings/hyena_coarse/\n",
      "entity_db/type_mappings/hyena_coarse/config.json\n",
      "entity_db/type_mappings/hyena_coarse/qid2typenames.json\n",
      "entity_db/type_mappings/hyena_coarse/qid2typeids.json\n",
      "entity_db/type_mappings/hyena_coarse/type_vocab.json\n",
      "entity_db/type_mappings/wiki/\n",
      "entity_db/type_mappings/wiki/qid2typeids.json\n",
      "entity_db/type_mappings/wiki/qid2typenames.json\n",
      "entity_db/type_mappings/wiki/type_vocab.json\n",
      "entity_db/type_mappings/wiki/config.json\n",
      "entity_db/type_mappings/wiki/type_vocab_to_wikidataqid.json\n"
     ]
    }
   ],
   "source": [
    "!sh download_model.sh uncased\n",
    "!sh download_data.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "\n",
    "# set up logging\n",
    "import sys\n",
    "import logging\n",
    "from importlib import reload\n",
    "\n",
    "reload(logging)\n",
    "# Set to logging.DEBUG for more logging output\n",
    "logging.basicConfig(\n",
    "    stream=sys.stdout, format=\"%(asctime)s %(message)s\", level=logging.INFO\n",
    ")\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# root_dir = FILL IN FULL PATH TO DIRECTORY WHERE DATA IS DOWNLOADED (i.e., root_dir/data and root_dir/models)\n",
    "root_dir = Path(\".\")\n",
    "# entity_dir = FILL IN PATH TO ENTITY_DB DATA (i.e., tutorial_data/data\n",
    "data_dir = root_dir / \"data\"\n",
    "entity_dir = data_dir / \"entity_db\"\n",
    "cand_map = entity_dir / \"entity_mappings/alias2qids.json\"\n",
    "# model_dir = FILL IN PATH TO MODELS\n",
    "model_dir = root_dir / \"models\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have a GPU with at least 12GB of memory available, set the below to 0 to run inference on a GPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = -1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Detect Mentions\n",
    "Bootleg uses a simple mention extraction algorithm that extracts mentions using a given candidate map. We will use a Wikipedia candidate map that we mined using Wikipedia anchor links and Wikidata aliases for a total of ~15 million mentions (provided in the Requirements section of this notebook)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the input dataset for the end-to-end pipeline, we assume a jsonlines file with a single dictionary with the key \"sentence\" and value as the text of the sentence, per line. For instance, you may have a file with the lines:\n",
    "\n",
    "    {\"sentence\": \"Who did the voice of the magician in Frosty the Snowman\"}\n",
    "    {\"sentence\": \"What is considered the Outer Banks in North Carolina\"}\n",
    "    \n",
    "Below, we have additional keys to keep track of the hand-labelled mentions, but this is purely for evaluating the quality of the end-to-end pipeline and is not needed in the common use cases of using Bootleg to detect and label mentions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "nq_sample_orig = data_dir / \"nq/test_50.jsonl\"\n",
    "nq_sample_bootleg = data_dir / \"nq/test_50_bootleg.jsonl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from bootleg.end2end.extract_mentions import extract_mentions\n",
    "\n",
    "verbose = False\n",
    "extract_mentions(\n",
    "    in_filepath=nq_sample_orig,\n",
    "    out_filepath=nq_sample_bootleg,\n",
    "    entity_db_dir=entity_dir,\n",
    "    verbose=verbose,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By looking at a sample of the extracted mentions, we can compare the mention extraction phase to the hand-labelled mentions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sentence</th>\n",
       "      <th>aliases_hand</th>\n",
       "      <th>spans_hand</th>\n",
       "      <th>aliases_bootleg</th>\n",
       "      <th>spans_bootleg</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Who did the voice of the magician in Frosty the Snowman</td>\n",
       "      <td>[frosty the snowman]</td>\n",
       "      <td>[[8, 11]]</td>\n",
       "      <td>[voice of, magician, frosty the snowman]</td>\n",
       "      <td>[[3, 5], [6, 7], [8, 11]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>What is considered the Outer Banks in North Carolina</td>\n",
       "      <td>[outer banks, north carolina]</td>\n",
       "      <td>[[4, 6], [7, 9]]</td>\n",
       "      <td>[outer banks, north carolina]</td>\n",
       "      <td>[[4, 6], [7, 9]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>The Nashville sound brought a polished and cosmopolitan sound to country music by</td>\n",
       "      <td>[nashville sound, country music]</td>\n",
       "      <td>[[1, 3], [10, 12]]</td>\n",
       "      <td>[nashville sound, music by]</td>\n",
       "      <td>[[1, 3], [11, 13]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What channel is the Premier League on in France</td>\n",
       "      <td>[premier league, france]</td>\n",
       "      <td>[[4, 6], [8, 9]]</td>\n",
       "      <td>[premier league, france]</td>\n",
       "      <td>[[4, 6], [8, 9]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>I Love It ( feat . Charli XCX ) Icona Pop</td>\n",
       "      <td>[i love it, charli xcx, icona pop]</td>\n",
       "      <td>[[0, 3], [6, 8], [9, 11]]</td>\n",
       "      <td>[charli xcx, icona pop]</td>\n",
       "      <td>[[6, 8], [9, 11]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>The U.S. Supreme Court hears appeals from circuit courts</td>\n",
       "      <td>[u.s. supreme court, circuit courts]</td>\n",
       "      <td>[[1, 4], [7, 9]]</td>\n",
       "      <td>[us supreme court, circuit courts]</td>\n",
       "      <td>[[1, 4], [7, 9]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Why does the author say that the vampire in Nosferatu is named Count Orlok and not Count Dracula</td>\n",
       "      <td>[nosferatu, count orlok, count dracula]</td>\n",
       "      <td>[[9, 10], [12, 14], [16, 18]]</td>\n",
       "      <td>[vampire, nosferatu, count orlok, count dracula]</td>\n",
       "      <td>[[7, 8], [9, 10], [12, 14], [16, 18]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Is there an active volcano in New Zealand</td>\n",
       "      <td>[new zealand]</td>\n",
       "      <td>[[6, 8]]</td>\n",
       "      <td>[volcano, new zealand]</td>\n",
       "      <td>[[4, 5], [6, 8]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Once Upon a Time Season 6 episode list</td>\n",
       "      <td>[once upon a time season 6]</td>\n",
       "      <td>[[0, 6]]</td>\n",
       "      <td>[upon a time, season 6, episode list]</td>\n",
       "      <td>[[1, 4], [4, 6], [6, 8]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Who is the former co-chairman Goldman Sachs who became a U.S. Secretary of the Treasury</td>\n",
       "      <td>[goldman sachs, us secretary of the treasury]</td>\n",
       "      <td>[[5, 7], [10, 15]]</td>\n",
       "      <td>[goldman sachs, us secretary of the treasury]</td>\n",
       "      <td>[[5, 7], [10, 15]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Who plays Norman Bates in the TV show</td>\n",
       "      <td>[norman bates]</td>\n",
       "      <td>[[2, 4]]</td>\n",
       "      <td>[norman bates]</td>\n",
       "      <td>[[2, 4]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Hitchhiker 's Guide to the Galaxy Slartibartfast quotes</td>\n",
       "      <td>[hitchhiker 's guide to the galaxy, slartibartfast]</td>\n",
       "      <td>[[0, 6], [6, 7]]</td>\n",
       "      <td>[hitchhiker s guide to the galaxy]</td>\n",
       "      <td>[[0, 6]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>What was Dennis Hopper 's bike in Easy Rider</td>\n",
       "      <td>[dennis hopper, easy rider]</td>\n",
       "      <td>[[2, 4], [7, 9]]</td>\n",
       "      <td>[dennis hopper, easy rider]</td>\n",
       "      <td>[[2, 4], [7, 9]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Where was 10 Things I Hate About You filmed school</td>\n",
       "      <td>[10 things i hate about you]</td>\n",
       "      <td>[[2, 8]]</td>\n",
       "      <td>[10 things i hate about you]</td>\n",
       "      <td>[[2, 8]]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Where does the last name Aponte come from</td>\n",
       "      <td>[aponte]</td>\n",
       "      <td>[[5, 6]]</td>\n",
       "      <td>[aponte]</td>\n",
       "      <td>[[5, 6]]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                                            sentence  \\\n",
       "0                                            Who did the voice of the magician in Frosty the Snowman   \n",
       "1                                               What is considered the Outer Banks in North Carolina   \n",
       "2                  The Nashville sound brought a polished and cosmopolitan sound to country music by   \n",
       "3                                                    What channel is the Premier League on in France   \n",
       "4                                                          I Love It ( feat . Charli XCX ) Icona Pop   \n",
       "5                                           The U.S. Supreme Court hears appeals from circuit courts   \n",
       "6   Why does the author say that the vampire in Nosferatu is named Count Orlok and not Count Dracula   \n",
       "7                                                          Is there an active volcano in New Zealand   \n",
       "8                                                             Once Upon a Time Season 6 episode list   \n",
       "9            Who is the former co-chairman Goldman Sachs who became a U.S. Secretary of the Treasury   \n",
       "10                                                             Who plays Norman Bates in the TV show   \n",
       "11                                           Hitchhiker 's Guide to the Galaxy Slartibartfast quotes   \n",
       "12                                                      What was Dennis Hopper 's bike in Easy Rider   \n",
       "13                                                Where was 10 Things I Hate About You filmed school   \n",
       "14                                                         Where does the last name Aponte come from   \n",
       "\n",
       "                                           aliases_hand  \\\n",
       "0                                  [frosty the snowman]   \n",
       "1                         [outer banks, north carolina]   \n",
       "2                      [nashville sound, country music]   \n",
       "3                              [premier league, france]   \n",
       "4                    [i love it, charli xcx, icona pop]   \n",
       "5                  [u.s. supreme court, circuit courts]   \n",
       "6               [nosferatu, count orlok, count dracula]   \n",
       "7                                         [new zealand]   \n",
       "8                           [once upon a time season 6]   \n",
       "9         [goldman sachs, us secretary of the treasury]   \n",
       "10                                       [norman bates]   \n",
       "11  [hitchhiker 's guide to the galaxy, slartibartfast]   \n",
       "12                          [dennis hopper, easy rider]   \n",
       "13                         [10 things i hate about you]   \n",
       "14                                             [aponte]   \n",
       "\n",
       "                       spans_hand  \\\n",
       "0                       [[8, 11]]   \n",
       "1                [[4, 6], [7, 9]]   \n",
       "2              [[1, 3], [10, 12]]   \n",
       "3                [[4, 6], [8, 9]]   \n",
       "4       [[0, 3], [6, 8], [9, 11]]   \n",
       "5                [[1, 4], [7, 9]]   \n",
       "6   [[9, 10], [12, 14], [16, 18]]   \n",
       "7                        [[6, 8]]   \n",
       "8                        [[0, 6]]   \n",
       "9              [[5, 7], [10, 15]]   \n",
       "10                       [[2, 4]]   \n",
       "11               [[0, 6], [6, 7]]   \n",
       "12               [[2, 4], [7, 9]]   \n",
       "13                       [[2, 8]]   \n",
       "14                       [[5, 6]]   \n",
       "\n",
       "                                     aliases_bootleg  \\\n",
       "0           [voice of, magician, frosty the snowman]   \n",
       "1                      [outer banks, north carolina]   \n",
       "2                        [nashville sound, music by]   \n",
       "3                           [premier league, france]   \n",
       "4                            [charli xcx, icona pop]   \n",
       "5                 [us supreme court, circuit courts]   \n",
       "6   [vampire, nosferatu, count orlok, count dracula]   \n",
       "7                             [volcano, new zealand]   \n",
       "8              [upon a time, season 6, episode list]   \n",
       "9      [goldman sachs, us secretary of the treasury]   \n",
       "10                                    [norman bates]   \n",
       "11                [hitchhiker s guide to the galaxy]   \n",
       "12                       [dennis hopper, easy rider]   \n",
       "13                      [10 things i hate about you]   \n",
       "14                                          [aponte]   \n",
       "\n",
       "                            spans_bootleg  \n",
       "0               [[3, 5], [6, 7], [8, 11]]  \n",
       "1                        [[4, 6], [7, 9]]  \n",
       "2                      [[1, 3], [11, 13]]  \n",
       "3                        [[4, 6], [8, 9]]  \n",
       "4                       [[6, 8], [9, 11]]  \n",
       "5                        [[1, 4], [7, 9]]  \n",
       "6   [[7, 8], [9, 10], [12, 14], [16, 18]]  \n",
       "7                        [[4, 5], [6, 8]]  \n",
       "8                [[1, 4], [4, 6], [6, 8]]  \n",
       "9                      [[5, 7], [10, 15]]  \n",
       "10                               [[2, 4]]  \n",
       "11                               [[0, 6]]  \n",
       "12                       [[2, 4], [7, 9]]  \n",
       "13                               [[2, 8]]  \n",
       "14                               [[5, 6]]  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from utils import load_mentions\n",
    "\n",
    "orig_mentions_df = load_mentions(nq_sample_orig)\n",
    "bootleg_mentions_df = load_mentions(nq_sample_bootleg)\n",
    "\n",
    "# join dataframes and sample\n",
    "res = pd.merge(\n",
    "    orig_mentions_df,\n",
    "    bootleg_mentions_df,\n",
    "    on=[\"sentence\"],\n",
    "    suffixes=[\"_hand\", \"_bootleg\"],\n",
    ")\n",
    "display(res.head(15))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the sample above, we see that generally Bootleg detects the same mentions as the hand-labelled mentions, however sometimes Bootleg extracts extra mentions (e.g \"colonies\" in \"Where did Britain create colonies for its empire\"). This is expected as we would rather the mention detection step filter out too few mentions than too many. It will be the job of the backbone model and postprocessing to filter out these extra mentions, by either thresholding the prediction probability or predicting a candidate that represents \"No Candidate\" (we refer to this as \"NC\").  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Disambiguate Mentions to Entities\n",
    "\n",
    "We run inference using a pretrained Bootleg model to disambiguate the extracted mentions to Wikidata QIDs. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, load the model config so we can set additional parameters and load the saved model during evaluation. We need to update the config parameters to point to the downloaded model checkpoint and data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/setuptools/distutils_patch.py:25: UserWarning: Distutils was imported before Setuptools. This usage is discouraged and may exhibit undesirable behaviors or errors. Please use Setuptools' objects directly or at least import Setuptools first.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from bootleg.utils.parser.parser_utils import parse_boot_and_emm_args\n",
    "from bootleg.utils.utils import load_yaml_file\n",
    "from bootleg.run import run_model\n",
    "\n",
    "config_in_path = model_dir / \"bootleg_uncased/bootleg_config.yaml\"\n",
    "\n",
    "config_args = load_yaml_file(config_in_path)\n",
    "\n",
    "# decrease number of data threads as this is a small file\n",
    "config_args[\"run_config\"][\"dataset_threads\"] = 2\n",
    "config_args[\"run_config\"][\"log_level\"] = \"info\"\n",
    "# set the model checkpoint path\n",
    "config_args[\"emmental\"][\"model_path\"] = str(\n",
    "    model_dir / \"bootleg_uncased/bootleg_wiki.pth\"\n",
    ")\n",
    "\n",
    "# set the path for the entity db and candidate map\n",
    "config_args[\"data_config\"][\"entity_dir\"] = str(entity_dir)\n",
    "config_args[\"data_config\"][\"alias_cand_map\"] = \"alias2qids.json\"\n",
    "\n",
    "# set the data path and kore50 test file\n",
    "config_args[\"data_config\"][\"data_dir\"] = str(data_dir / \"nq\")\n",
    "\n",
    "# to speed things up for the tutorial, we have already prepped the data with the mentions detected by Bootleg\n",
    "config_args[\"data_config\"][\"test_dataset\"][\"file\"] = nq_sample_bootleg.name\n",
    "\n",
    "# set the devie if on CPU\n",
    "config_args[\"emmental\"][\"device\"] = device\n",
    "\n",
    "# save the new args (helps if you want to run things via command line)\n",
    "config_args = parse_boot_and_emm_args(config_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-10-15 22:29:58,410 Setting logging directory to: bootleg-logs/bootleg_wiki\n",
      "2021-10-15 22:29:58,453 Loading Emmental default config from /lfs/raiders3/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/emmental-default-config.yaml.\n",
      "2021-10-15 22:29:58,454 Updating Emmental config from user provided config.\n",
      "2021-10-15 22:29:58,455 Set random seed to 1234.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[10/15/21 22:29:58] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> COMMAND: <span style=\"color: #800080; text-decoration-color: #800080\">/lfs/local/0/senwu/.pyenv/versions/3.8.6</span> <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             <span style=\"color: #800080; text-decoration-color: #800080\">/envs/venv38/lib/python3.8/site-packages/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">ipykerne</span>               \n",
       "                             <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">l_launcher.py</span> -f <span style=\"color: #800080; text-decoration-color: #800080\">/lfs/raiders3/0/senwu/.local/sha</span>               \n",
       "                             <span style=\"color: #800080; text-decoration-color: #800080\">re/jupyter/runtime/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">kernel-b72e10ac-199d-41f8-af7a</span>               \n",
       "                             <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">-6a3faa67756a.json</span>                                              \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/15/21 22:29:58]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m COMMAND: \u001b[35m/lfs/local/0/senwu/.pyenv/versions/3.8.6\u001b[0m \u001b]8;id=1634362198.6277988-462141;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             \u001b[35m/envs/venv38/lib/python3.8/site-packages/\u001b[0m\u001b[95mipykerne\u001b[0m               \n",
       "                             \u001b[95ml_launcher.py\u001b[0m -f \u001b[35m/lfs/raiders3/0/senwu/.local/sha\u001b[0m               \n",
       "                             \u001b[35mre/jupyter/runtime/\u001b[0m\u001b[95mkernel-b72e10ac-199d-41f8-af7a\u001b[0m               \n",
       "                             \u001b[95m-6a3faa67756a.json\u001b[0m                                              \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Saving config to                                  <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             bootleg-logs/bootleg_wiki/parsed_config.yaml                    \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Saving config to                                  \u001b]8;id=1634362198.6319675-122526;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             bootleg-logs/bootleg_wiki/parsed_config.yaml                    \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Git Hash: <span style=\"color: #008000; text-decoration-color: #008000\">b'06a9297-2021-10-15'</span>                   <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Git Hash: \u001b[32mb'06a9297-2021-10-15'\u001b[0m                   \u001b]8;id=1634362198.8429523-7855;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loading entity symbols<span style=\"color: #808000; text-decoration-color: #808000\">...</span>                         <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading entity symbols\u001b[33m...\u001b[0m                         \u001b]8;id=1634362198.8474433-95046;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[10/15/21 22:32:48] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Starting to build data for test from              <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             data/nq/test_50_bootleg.jsonl                                   \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/15/21 22:32:48]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Starting to build data for test from              \u001b]8;id=1634362368.2099764-955227;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             data/nq/test_50_bootleg.jsonl                                   \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/bootleg/bootleg/dataset.py:1283: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  guid_dtype = np.dtype(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loading data from data/nq/prep/test_50_bootleg_be <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             rt-base-uncased_L128_E128_W64_T<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_K<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_D<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_I               \n",
       "                             nC1_Aug1/ned_data.bin and data/nq/prep/test_50_bo               \n",
       "                             otleg_bert-base-uncased_L128_E128_W64_T<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_K<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>               \n",
       "                             _D<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_InC1_Aug1/ned_label.bin                                  \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading data from data/nq/prep/test_50_bootleg_be \u001b]8;id=1634362368.790823-988319;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             rt-base-uncased_L128_E128_W64_T\u001b[3;92mTrue\u001b[0m_K\u001b[3;92mTrue\u001b[0m_D\u001b[3;92mTrue\u001b[0m_I               \n",
       "                             nC1_Aug1/ned_data.bin and data/nq/prep/test_50_bo               \n",
       "                             otleg_bert-base-uncased_L128_E128_W64_T\u001b[3;92mTrue\u001b[0m_K\u001b[3;92mTrue\u001b[0m               \n",
       "                             _D\u001b[3;92mTrue\u001b[0m_InC1_Aug1/ned_label.bin                                  \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/numpy/core/memmap.py:230: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  descr = dtypedescr(dtype)\n",
      "/bootleg/bootleg/dataset.py:1561: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  ../torch/csrc/utils/tensor_numpy.cpp:180.)\n",
      "  X_dict[\"sent_idx\"] = torch.from_numpy(mmap_file[\"sent_idx\"])\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Final data initialization time for test is        <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.</span>6658532619476318s                                             \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Final data initialization time for test is        \u001b]8;id=1634362368.881055-984894;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             \u001b[1;36m0.\u001b[0m6658532619476318s                                             \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Built dataloader for test set with <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">89</span> and <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>       <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             threads samples <span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">Shuffle</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>, Batch <span style=\"color: #808000; text-decoration-color: #808000\">size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span><span style=\"font-weight: bold\">)</span>.                 \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Built dataloader for test set with \u001b[1;36m89\u001b[0m and \u001b[1;36m2\u001b[0m       \u001b]8;id=1634362368.962654-846203;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             threads samples \u001b[1m(\u001b[0m\u001b[33mShuffle\u001b[0m=\u001b[3;91mFalse\u001b[0m, Batch \u001b[33msize\u001b[0m=\u001b[1;36m32\u001b[0m\u001b[1m)\u001b[0m.                 \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Building slice dataset for test from              <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             data/nq/test_50_bootleg.jsonl.                                  \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Building slice dataset for test from              \u001b]8;id=1634362368.9666245-610509;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             data/nq/test_50_bootleg.jsonl.                                  \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[10/15/21 22:32:49] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loading data from data/nq/prep/test_50_bootleg_be <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             rt-base-uncased_L128_E128_W64_T<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_K<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_D<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_I               \n",
       "                             nC1_Aug1/ned_slices_6a0e90c2ae.bin and data/nq/pr               \n",
       "                             ep/test_50_bootleg_bert-base-uncased_L128_E128_W6               \n",
       "                             4_T<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_K<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_D<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>_InC1_Aug1/ned_slices_config.j               \n",
       "                             son                                                             \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/15/21 22:32:49]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading data from data/nq/prep/test_50_bootleg_be \u001b]8;id=1634362369.091449-36622;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             rt-base-uncased_L128_E128_W64_T\u001b[3;92mTrue\u001b[0m_K\u001b[3;92mTrue\u001b[0m_D\u001b[3;92mTrue\u001b[0m_I               \n",
       "                             nC1_Aug1/ned_slices_6a0e90c2ae.bin and data/nq/pr               \n",
       "                             ep/test_50_bootleg_bert-base-uncased_L128_E128_W6               \n",
       "                             4_T\u001b[3;92mTrue\u001b[0m_K\u001b[3;92mTrue\u001b[0m_D\u001b[3;92mTrue\u001b[0m_InC1_Aug1/ned_slices_config.j               \n",
       "                             son                                                             \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building sent idx to row idx mapping: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:00<00:00, 10668.73it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Final slice data initialization time from test is <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.</span>2833065986633301s                                             \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Final slice data initialization time from test is \u001b]8;id=1634362369.2500427-704185;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             \u001b[1;36m0.\u001b[0m2833065986633301s                                             \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Starting Bootleg Model                            <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Starting Bootleg Model                            \u001b]8;id=1634362369.2536874-726754;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Created emmental model Bootleg that contains task    <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">model.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:72</span>\n",
       "                             set<span style=\"font-weight: bold\">()</span>.                                                          \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Created emmental model Bootleg that contains task    \u001b]8;id=1634362369.2579405-103202;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\u001b\\\u001b[2mmodel.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:72\u001b[0m\n",
       "                             set\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m.                                                          \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.decoder.weight', 'cls.seq_relationship.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.decoder.weight', 'cls.seq_relationship.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[10/15/21 22:33:23] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Created task: NED                                     <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/task.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">task.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:78</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/15/21 22:33:23]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Created task: NED                                     \u001b]8;id=1634362403.2959192-371862;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/task.py\u001b\\\u001b[2mtask.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:78\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Moving context_encoder module to CPU.               <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">model.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:108</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Moving context_encoder module to CPU.               \u001b]8;id=1634362403.3012805-248313;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\u001b\\\u001b[2mmodel.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:108\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Moving entity_encoder module to CPU.                <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">model.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:108</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Moving entity_encoder module to CPU.                \u001b]8;id=1634362403.3076296-18274;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\u001b\\\u001b[2mmodel.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:108\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> <span style=\"font-weight: bold\">[</span>Bootleg<span style=\"font-weight: bold\">]</span> Model loaded from                         <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">model.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:877</span>\n",
       "                             models/bootleg_uncased/bootleg_wiki.pth                         \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m \u001b[1m[\u001b[0mBootleg\u001b[1m]\u001b[0m Model loaded from                         \u001b]8;id=1634362403.7278533-32310;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\u001b\\\u001b[2mmodel.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:877\u001b[0m\n",
       "                             models/bootleg_uncased/bootleg_wiki.pth                         \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Moving context_encoder module to CPU.               <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">model.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:108</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Moving context_encoder module to CPU.               \u001b]8;id=1634362403.731594-827088;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\u001b\\\u001b[2mmodel.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:108\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Moving entity_encoder module to CPU.                <a href=\"file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">model.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:108</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Moving entity_encoder module to CPU.                \u001b]8;id=1634362403.7381322-16609;file:///lfs/local/0/senwu/.pyenv/versions/3.8.6/envs/venv38/lib/python3.8/site-packages/emmental/model.py\u001b\\\u001b[2mmodel.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:108\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating Bootleg (test): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:56<00:00, 18.96s/it]\n",
      "/bootleg/bootleg/utils/eval_utils.py:600: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  unmerged_storage_type = np.dtype(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[10/15/21 22:34:23] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Finished dumping. Merging results across          <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             accumulation steps.                                             \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/15/21 22:34:23]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Finished dumping. Merging results across          \u001b]8;id=1634362463.3506913-362900;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             accumulation steps.                                             \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Bootleg labels saved at bootleg-logs/bootleg_wiki <a href=\"file:///bootleg/bootleg/__init__.py\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">__init__.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:11</span>\n",
       "                             /test_50_bootleg/bootleg_wiki/bootleg_labels.json               \n",
       "                             l                                                               \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Bootleg labels saved at bootleg-logs/bootleg_wiki \u001b]8;id=1634362463.4392364-679808;file:///bootleg/bootleg/__init__.py\u001b\\\u001b[2m__init__.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:11\u001b[0m\n",
       "                             /test_50_bootleg/bootleg_wiki/bootleg_labels.json               \n",
       "                             l                                                               \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bootleg_label_file = run_model(mode=\"dump_preds\", config=config_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "**Note that Bootleg automatically handles prepping of new data files for running. These are all saved in `data_config.data_dir`/`data_config.data_prep_dir`. If you change the contents of the underlying `jsonl` file _without_ removing the saved prep file or setting `data_config.overwrite_preprocessed_data` to be `True`, Bootleg will reuse the old prepped file.**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now evaluate the overall quality of the end-to-end pipeline via precision / recall metrics, where the *recall* indicates what proportion of the hand-labelled mentions Bootleg correctly detects and disambiguates, and *precision* indicates what proportion of the mentions that Bootleg labels are correct. For instance, if Bootleg only labelled the few mentions it was very confident in, then it would have a low recall and high precision.\n",
    "\n",
    "To detect if mentions match the hand-labelled mention spans, we report weak and exact match metrics. Weak means the predicted and gold span boundaries just need to overlap for an entity (e.g., predicted mention 'the wizard of oz' is counted as correct for the gold mention 'wizard of oz' if the correct entity is predicted). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WEAK MATCHING\n",
      "precision = 62 / 79 = 0.7848101265822784\n",
      "recall = 62 / 78 = 0.7948717948717948\n",
      "f1 = 0.7898089171974522\n",
      "\n",
      "EXACT MATCHING\n",
      "precision = 62 / 79 = 0.7848101265822784\n",
      "recall = 62 / 78 = 0.7948717948717948\n",
      "f1 = 0.7898089171974522\n"
     ]
    }
   ],
   "source": [
    "from utils import compute_metrics\n",
    "\n",
    "bootleg_end2end_errors = compute_metrics(\n",
    "    gold_file=nq_sample_orig, pred_file=bootleg_label_file, threshold=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can examine errors in the end-to-end pipeline below. As you increase the threshold in the `compute_metrics` command, entities with a prediction probability less than the threshold will be filtered out. If too few entities are predicted, lowering the threshold may help.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sent_idx</th>\n",
       "      <th>text</th>\n",
       "      <th>gold_aliases</th>\n",
       "      <th>gold_qids</th>\n",
       "      <th>gold_spans</th>\n",
       "      <th>pred_aliases</th>\n",
       "      <th>pred_qids</th>\n",
       "      <th>pred_spans</th>\n",
       "      <th>pred_probs</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>16</td>\n",
       "      <td>Where did Britain create colonies for its empire</td>\n",
       "      <td>[britain, empire]</td>\n",
       "      <td>[Q161885, Q8680]</td>\n",
       "      <td>[[2, 3], [7, 8]]</td>\n",
       "      <td>[britain]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>34</td>\n",
       "      <td>Who 's doing the halftime show in 2018</td>\n",
       "      <td>[halftime show]</td>\n",
       "      <td>[Q902899]</td>\n",
       "      <td>[[4, 6]]</td>\n",
       "      <td>[halftime show]</td>\n",
       "      <td>[Q11161626]</td>\n",
       "      <td>[[4, 6]]</td>\n",
       "      <td>[0.8461166620254517]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4</td>\n",
       "      <td>I Love It ( feat . Charli XCX ) Icona Pop</td>\n",
       "      <td>[i love it, charli xcx, icona pop]</td>\n",
       "      <td>[Q3273659, Q5084390, Q808703]</td>\n",
       "      <td>[[0, 3], [6, 8], [9, 11]]</td>\n",
       "      <td>[charli xcx, icona pop]</td>\n",
       "      <td>[Q5084390, Q808703]</td>\n",
       "      <td>[[6, 8], [9, 11]]</td>\n",
       "      <td>[1.0, 0.9958118200302124]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>Is there an active volcano in New Zealand</td>\n",
       "      <td>[new zealand]</td>\n",
       "      <td>[Q664]</td>\n",
       "      <td>[[6, 8]]</td>\n",
       "      <td>[volcano, new zealand]</td>\n",
       "      <td>[Q8072, Q664]</td>\n",
       "      <td>[[4, 5], [6, 8]]</td>\n",
       "      <td>[0.9965059757232666, 0.9999643564224243]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>The Nashville sound brought a polished and cosmopolitan sound to country music by</td>\n",
       "      <td>[nashville sound, country music]</td>\n",
       "      <td>[Q1751782, Q83440]</td>\n",
       "      <td>[[1, 3], [10, 12]]</td>\n",
       "      <td>[nashville sound, music by]</td>\n",
       "      <td>[Q1751782, Q17059875]</td>\n",
       "      <td>[[1, 3], [11, 13]]</td>\n",
       "      <td>[0.9996482133865356, 1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>40</td>\n",
       "      <td>Where does the last name Vigil come from</td>\n",
       "      <td>[vigil]</td>\n",
       "      <td>[Q16878937]</td>\n",
       "      <td>[[5, 6]]</td>\n",
       "      <td>[vigil]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>11</td>\n",
       "      <td>Hitchhiker 's Guide to the Galaxy Slartibartfast quotes</td>\n",
       "      <td>[hitchhiker 's guide to the galaxy, slartibartfast]</td>\n",
       "      <td>[Q25169, Q779920]</td>\n",
       "      <td>[[0, 6], [6, 7]]</td>\n",
       "      <td>[hitchhiker s guide to the galaxy]</td>\n",
       "      <td>[Q25169]</td>\n",
       "      <td>[[0, 6]]</td>\n",
       "      <td>[1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>8</td>\n",
       "      <td>Once Upon a Time Season 6 episode list</td>\n",
       "      <td>[once upon a time season 6]</td>\n",
       "      <td>[Q23301616]</td>\n",
       "      <td>[[0, 6]]</td>\n",
       "      <td>[upon a time, season 6, episode list]</td>\n",
       "      <td>[Q55636748]</td>\n",
       "      <td>[[1, 4]]</td>\n",
       "      <td>[1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>Who did the voice of the magician in Frosty the Snowman</td>\n",
       "      <td>[frosty the snowman]</td>\n",
       "      <td>[Q5506238]</td>\n",
       "      <td>[[8, 11]]</td>\n",
       "      <td>[voice of, magician, frosty the snowman]</td>\n",
       "      <td>[Q11540682, Q5506238]</td>\n",
       "      <td>[[3, 5], [8, 11]]</td>\n",
       "      <td>[1.0, 0.9153743982315063]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>19</td>\n",
       "      <td>Who played the bank robber in Dirty Harry</td>\n",
       "      <td>[dirty harry]</td>\n",
       "      <td>[Q110206]</td>\n",
       "      <td>[[6, 8]]</td>\n",
       "      <td>[bank robber, dirty harry]</td>\n",
       "      <td>[Q806824, Q110206]</td>\n",
       "      <td>[[3, 5], [6, 8]]</td>\n",
       "      <td>[0.7667599320411682, 0.6623506546020508]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    sent_idx  \\\n",
       "7         16   \n",
       "13        34   \n",
       "2          4   \n",
       "3          7   \n",
       "1          2   \n",
       "16        40   \n",
       "5         11   \n",
       "4          8   \n",
       "0          0   \n",
       "8         19   \n",
       "\n",
       "                                                                                 text  \\\n",
       "7                                    Where did Britain create colonies for its empire   \n",
       "13                                             Who 's doing the halftime show in 2018   \n",
       "2                                           I Love It ( feat . Charli XCX ) Icona Pop   \n",
       "3                                           Is there an active volcano in New Zealand   \n",
       "1   The Nashville sound brought a polished and cosmopolitan sound to country music by   \n",
       "16                                           Where does the last name Vigil come from   \n",
       "5                             Hitchhiker 's Guide to the Galaxy Slartibartfast quotes   \n",
       "4                                              Once Upon a Time Season 6 episode list   \n",
       "0                             Who did the voice of the magician in Frosty the Snowman   \n",
       "8                                           Who played the bank robber in Dirty Harry   \n",
       "\n",
       "                                           gold_aliases  \\\n",
       "7                                     [britain, empire]   \n",
       "13                                      [halftime show]   \n",
       "2                    [i love it, charli xcx, icona pop]   \n",
       "3                                         [new zealand]   \n",
       "1                      [nashville sound, country music]   \n",
       "16                                              [vigil]   \n",
       "5   [hitchhiker 's guide to the galaxy, slartibartfast]   \n",
       "4                           [once upon a time season 6]   \n",
       "0                                  [frosty the snowman]   \n",
       "8                                         [dirty harry]   \n",
       "\n",
       "                        gold_qids                 gold_spans  \\\n",
       "7                [Q161885, Q8680]           [[2, 3], [7, 8]]   \n",
       "13                      [Q902899]                   [[4, 6]]   \n",
       "2   [Q3273659, Q5084390, Q808703]  [[0, 3], [6, 8], [9, 11]]   \n",
       "3                          [Q664]                   [[6, 8]]   \n",
       "1              [Q1751782, Q83440]         [[1, 3], [10, 12]]   \n",
       "16                    [Q16878937]                   [[5, 6]]   \n",
       "5               [Q25169, Q779920]           [[0, 6], [6, 7]]   \n",
       "4                     [Q23301616]                   [[0, 6]]   \n",
       "0                      [Q5506238]                  [[8, 11]]   \n",
       "8                       [Q110206]                   [[6, 8]]   \n",
       "\n",
       "                                pred_aliases              pred_qids  \\\n",
       "7                                  [britain]                     []   \n",
       "13                           [halftime show]            [Q11161626]   \n",
       "2                    [charli xcx, icona pop]    [Q5084390, Q808703]   \n",
       "3                     [volcano, new zealand]          [Q8072, Q664]   \n",
       "1                [nashville sound, music by]  [Q1751782, Q17059875]   \n",
       "16                                   [vigil]                     []   \n",
       "5         [hitchhiker s guide to the galaxy]               [Q25169]   \n",
       "4      [upon a time, season 6, episode list]            [Q55636748]   \n",
       "0   [voice of, magician, frosty the snowman]  [Q11540682, Q5506238]   \n",
       "8                 [bank robber, dirty harry]     [Q806824, Q110206]   \n",
       "\n",
       "            pred_spans                                pred_probs  \n",
       "7                   []                                        []  \n",
       "13            [[4, 6]]                      [0.8461166620254517]  \n",
       "2    [[6, 8], [9, 11]]                 [1.0, 0.9958118200302124]  \n",
       "3     [[4, 5], [6, 8]]  [0.9965059757232666, 0.9999643564224243]  \n",
       "1   [[1, 3], [11, 13]]                 [0.9996482133865356, 1.0]  \n",
       "16                  []                                        []  \n",
       "5             [[0, 6]]                                     [1.0]  \n",
       "4             [[1, 4]]                                     [1.0]  \n",
       "0    [[3, 5], [8, 11]]                 [1.0, 0.9153743982315063]  \n",
       "8     [[3, 5], [6, 8]]  [0.7667599320411682, 0.6623506546020508]  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(bootleg_end2end_errors).sample(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some of the errors Bootleg makes is predicting too general of a candidate (e.g. Oregon State Beavers instead of Oregon State Beavers baseball). Other errors are due to ambiguous sentences (e.g. \"cast of characters in fiddler on the roof\" -> should this be the movie or the musical?). Finally another bucket of errors suggests that we need to boost certain training signals -- this is an area we're actively pursuing in Bootleg with an investigation of model guidability!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}