tutorials/entity_embedding_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creates Entity Embeddings\n",
    "\n",
    "We will use the same requirements as in the end2end_ned_tutorial.ipynb file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>.container { width:90% !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:90% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2021-09-06 22:33:39,371][INFO] emmental.meta:122 - Setting logging directory to: /dfs/scratch0/lorr1/projects/bootleg/tutorial_data/data/log_dir\n",
      "[2021-09-06 22:33:39,428][INFO] emmental.meta:64 - Loading Emmental default config from /dfs/scratch0/lorr1/projects/emmental/src/emmental/emmental-default-config.yaml.\n",
      "[2021-09-06 22:33:39,429][INFO] emmental.meta:174 - Updating Emmental config from user provided config.\n",
      "[2021-09-06 22:33:39,431][INFO] emmental.utils.seed:27 - Set random seed to 1234.\n",
      "[2021-09-06 22:33:39,436][DEBUG] bootleg.end2end.bootleg_annotator:225 - Reading entity database\n",
      "[2021-09-06 22:35:38,197][DEBUG] bootleg.end2end.bootleg_annotator:238 - Reading entity database\n",
      "[2021-09-06 22:38:10,187][DEBUG] bootleg.end2end.bootleg_annotator:248 - Reading word tokenizers\n",
      "[2021-09-06 22:38:10,196][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:10,491][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/tokenizer_config.json HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:10,500][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:10,795][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/config.json HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:10,805][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:11,099][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/vocab.txt HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:11,108][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:11,401][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/tokenizer.json HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:11,408][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:11,698][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/added_tokens.json HTTP/1.1\" 404 0\n",
      "[2021-09-06 22:38:11,707][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:11,993][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/special_tokens_map.json HTTP/1.1\" 404 0\n",
      "[2021-09-06 22:38:12,000][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:12,293][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/tokenizer_config.json HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:12,355][INFO] emmental.model:72 - Created emmental model Bootleg that contains task set().\n",
      "[2021-09-06 22:38:12,360][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Standard Cands CrossEntropy Loss\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2021-09-06 22:38:12,647][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/config.json HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:12,657][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:12,954][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/pytorch_model.bin HTTP/1.1\" 302 0\n",
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.bias', 'cls.predictions.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",
      "[2021-09-06 22:38:14,773][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:15,073][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/config.json HTTP/1.1\" 200 0\n",
      "[2021-09-06 22:38:15,082][DEBUG] urllib3.connectionpool:971 - Starting new HTTPS connection (1): huggingface.co:443\n",
      "[2021-09-06 22:38:15,374][DEBUG] urllib3.connectionpool:452 - https://huggingface.co:443 \"HEAD /bert-base-uncased/resolve/main/pytorch_model.bin HTTP/1.1\" 302 0\n",
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.bias', 'cls.predictions.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",
      "[2021-09-06 22:38:17,067][INFO] emmental.task:78 - Created task: NED\n",
      "[2021-09-06 22:38:17,068][INFO] emmental.model:97 - Moving context_encoder module to GPU (cuda:0).\n",
      "[2021-09-06 22:38:23,086][INFO] emmental.model:97 - Moving entity_encoder module to GPU (cuda:0).\n",
      "[2021-09-06 22:38:23,271][DEBUG] bootleg.end2end.bootleg_annotator:270 - Loading model\n",
      "[2021-09-06 22:38:23,729][INFO] emmental.model:800 - [Bootleg] Model loaded from /dfs/scratch0/lorr1/projects/bootleg/tutorial_data/models/bootleg_uncased/bootleg_wiki.pth\n",
      "[2021-09-06 22:38:23,730][INFO] emmental.model:97 - Moving context_encoder module to GPU (cuda:0).\n",
      "[2021-09-06 22:38:23,732][INFO] emmental.model:97 - Moving entity_encoder module to GPU (cuda:0).\n",
      "[2021-09-06 22:38:23,736][DEBUG] bootleg.end2end.extract_mentions:108 - Loading candidate mapping...\n",
      "Reading candidate map: 100%|██████████| 15202497/15202497 [00:33<00:00, 447326.49it/s]\n",
      "[2021-09-06 22:38:57,724][DEBUG] bootleg.end2end.extract_mentions:111 - Loaded candidate mapping with 15202497 aliases.\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import sys\n",
    "import numpy as np\n",
    "import logging\n",
    "from importlib import reload\n",
    "reload(logging)\n",
    "\n",
    "reload(logging)\n",
    "# Set to logging.DEBUG for more logging output\n",
    "logging.basicConfig(stream=sys.stdout, format='%(asctime)s %(message)s', level=logging.INFO)\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",
    "# model_dir = FILL IN PATH TO MODELS\n",
    "model_dir = root_dir / \"models\"\n",
    "device = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "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.extract_all_entities 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",
    "# set the model checkpoint path\n",
    "config_args[\"emmental\"][\"model_path\"] = str(model_dir / 'bootleg_uncased/bootleg_wiki.pth')\n",
    "\n",
    "# set the path for the entity db and candidate map\n",
    "config_args[\"data_config\"][\"entity_dir\"] = str(entity_dir)\n",
    "\n",
    "# set the devie if on CPU\n",
    "config_args[\"emmental\"][\"device\"] = device\n",
    "\n",
    "# change dataloader threads to avoid memory issues with eval dataloader\n",
    "config_args[\"data_config\"][\"dataloader_threads\"] = 0\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)\n",
    "\n",
    "print(f\"Can also run\\npython -m bootleg.extract_all_entities --config_script {model_dir}/bootleg_uncased/bootleg_config.yaml \\\n",
    "    --emmental.model_path {model_dir}/bootleg_uncased/bootleg_wiki.pth --data_config.entity_dir {entity_dir} \\\n",
    "    --emmental.device {device} --data_config.dataloader_threads 0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "bootleg_embedding_file = run_model(config=config_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "embs = np.load(bootleg_embedding_file)\n",
    "print(embs.shape)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's see how to extact embeddings for specific entities"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from bootleg.symbols.entity_profile import EntityProfile\n",
    "\n",
    "ep = EntityProfile.load_from_cache(entity_dir)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "qid = \"Q21\"\n",
    "print(\"Title\", ep.get_title(qid))\n",
    "\n",
    "embedding = embs[ep.get_eid(qid)]\n",
    "\n",
    "import faiss                   # make faiss available\n",
    "index = faiss.IndexFlatL2(embs)\n",
    "\n",
    "D, I = index.search(embedding, 10)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}