tutorials/annotation-on-the-fly.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bootleg Annotator 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 on the fly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Requirements\n",
    "When evaluating Bootleg using the annotator, Bootleg processes possible mentions in text with three environment flags: ``BOOTLEG_STRIP``, ``BOOTLEG_LOWER``, ``BOOTLEG_LANG_CODE``. The first sets the language to use for Spacy. The second is if the user wants to strip punctuation on mentions (set to False by default). The third is if the user wants to call ``.lower()`` (set to True by default). Set these with `os.environ`.\n",
    "\n",
    "You will need to download the following files for this notebook:\n",
    "- Pretrained Bootleg uncased model and config [here](https://bootleg-ned-data.s3-us-west-1.amazonaws.com/models/latest/bootleg_uncased.tar.gz)\n",
    "- Entity data [here](https://bootleg-ned-data.s3-us-west-1.amazonaws.com/data/latest/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": [],
   "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",
    "\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",
    "# 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": [
    "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": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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.utils import load_yaml_file\n",
    "\n",
    "config_in_path = model_dir / \"bootleg_uncased/bootleg_config.yaml\"\n",
    "\n",
    "config_args = load_yaml_file(config_in_path)\n",
    "\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's give the config to load the annotator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-10-15 20:23:06,436 Setting logging directory to: bootleg-logs/bootleg_wiki\n",
      "2021-10-15 20:23:06,480 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 20:23:06,481 Updating Emmental config from user provided config.\n",
      "2021-10-15 20:23:06,482 Set random seed to 1234.\n",
      "2021-10-15 20:29:36,662 Created emmental model Bootleg that contains task set().\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight']\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.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight']\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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-10-15 20:29:46,048 Created task: NED\n",
      "2021-10-15 20:29:46,053 Moving context_encoder module to CPU.\n",
      "2021-10-15 20:29:46,057 Moving entity_encoder module to CPU.\n",
      "2021-10-15 20:29:46,746 [Bootleg] Model loaded from ../models/bootleg_uncased/bootleg_wiki.pth\n",
      "2021-10-15 20:29:46,747 Moving context_encoder module to CPU.\n",
      "2021-10-15 20:29:46,751 Moving entity_encoder module to CPU.\n"
     ]
    }
   ],
   "source": [
    "# Load new annotator with our config - notice how it does have to reprep some things\n",
    "from bootleg.end2end.bootleg_annotator import BootlegAnnotator\n",
    "\n",
    "# You can also pass `return_embs=True` to get the embeddings\n",
    "ann = BootlegAnnotator(\n",
    "    config=config_args, device=device, return_embs=False, verbose=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['Abraham Lincoln']]\n",
      "[['Lincoln Motor Company']]\n"
     ]
    }
   ],
   "source": [
    "print(ann.label_mentions([\"I am Lincoln\"])[\"titles\"])\n",
    "print(ann.label_mentions([\"How much is a Lincoln\"])[\"titles\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Faster inference\n",
    "If you want more efficient inference of the annotator, we have the ability for the user to pass in a static entity\n",
    "embedding matrix so the model does not have to call a forward pass of the entity encoder.\n",
    "\n",
    "See our ```entity_embedding_tutorial.ipynb``` for how to call ```extract_all_entities```. The output of this\n",
    "can be passed into our annotator via"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "entity_emb_file = \"<path to file>\"\n",
    "ann = BootlegAnnotator(config=config_args, device=device, return_embs=False, entity_emb_file=entity_emb_file)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}