NEL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('./src/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import sqlite3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from torchlite.torch.train_callbacks import ModelSaverCallback\n",
    "\n",
    "from model.embedding import ModelVectorizer, OnDiskVectorizer\n",
    "from model.arc2 import ARC2, PreConv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tokenizer(text, alpha_only=True):  # create a tokenizer function\n",
    "    words = [tok for tok in nltk.word_tokenize(text) if (not alpha_only or tok.isalpha())]\n",
    "    return words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NEL:\n",
    "    def __init__(self):\n",
    "        net_params = {\n",
    "            'preconv': True,\n",
    "            'word_emb_sizes': [300],\n",
    "            'preconv_size': [300],\n",
    "            'matrix_depth': [120],\n",
    "            'conv_depth': [120, 60, 60],\n",
    "            'out_size': [60]\n",
    "        }\n",
    "        self.vectorizer = OnDiskVectorizer(mtx_path='./data/fastText.nmy', meta_path='./data/fastText.json')\n",
    "        preconv = PreConv(\n",
    "            word_emb_sizes=net_params['word_emb_sizes'],\n",
    "            sent_conv_size=net_params['preconv_size'],\n",
    "            dropout=0.0,\n",
    "            window=2\n",
    "        )\n",
    "\n",
    "        self.model = ARC2(\n",
    "            vectorizer=None,\n",
    "            preconv=preconv,\n",
    "            matrix_depth=net_params['matrix_depth'],\n",
    "            conv_depth=net_params['conv_depth'],\n",
    "            out_size=net_params['out_size'],\n",
    "            window=2,\n",
    "            dropout=0.0\n",
    "        )\n",
    "        \n",
    "        ModelSaverCallback.restore_model_from_file(self.model, './data/models/ARC2_best.pth', load_with_cpu=True)\n",
    "        self.model = self.model.eval()\n",
    "        \n",
    "        self.connection = sqlite3.connect('./data/mentions.sqlite3')\n",
    "        self.cur = self.connection.cursor()\n",
    "        \n",
    "    def match_sentences(self, sent_a, sent_b):\n",
    "        sent_a_token = tokenizer(sent_a) + [' '] * 6  # ensure data size > conv kernel size\n",
    "        sent_b_token = tokenizer(sent_b) + [' '] * 6\n",
    "        \n",
    "        sent_a_vect = self.vectorizer.convert([sent_a_token])\n",
    "        sent_b_vect = self.vectorizer.convert([sent_b_token])\n",
    "        \n",
    "        score = self.model.forward(sent_a_vect, sent_b_vect)\n",
    "\n",
    "        return score[0][1].item()\n",
    "    \n",
    "    \n",
    "    def create_reference_set(self, mention):\n",
    "        self.cur.execute(\"select * from mentions where mention match ? limit 1000\", (mention, ))\n",
    "        \n",
    "        result = {}\n",
    "        \n",
    "        entities = defaultdict(list)\n",
    "        for entity, left_context, mention, right_context in nel.cur.fetchall():\n",
    "            entities[entity].append((left_context, mention, right_context))\n",
    "        for key, val in entities.items():\n",
    "            if len(val) > 2:\n",
    "                result[key] = list(map(lambda x: x[0] + \" XXXXX \" + x[2], random.sample(val, min(10, len(val)))))\n",
    "        \n",
    "        return result\n",
    "    \n",
    "    \n",
    "    def match_all(self, left_context, mention, right_context, sentences):\n",
    "        ref_sent = left_context + \" XXXXX \" + right_context\n",
    "        all_scores = []\n",
    "        for sent in sentences:\n",
    "            score = self.match_sentences(ref_sent, sent)\n",
    "            all_scores.append(score)\n",
    "            \n",
    "        return np.mean(sorted(all_scores, reverse=True)[:3])\n",
    "    \n",
    "    def disabiguate(self, left_context, mention, right_context):\n",
    "        train = nel.create_reference_set(mention)\n",
    "        \n",
    "        for entity, sents in train.items():\n",
    "            score = self.match_all(left_context, mention, right_context, sents)\n",
    "            print(entity, score)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Model restored ---\n",
      "\n",
      "CPU times: user 2.91 s, sys: 432 ms, total: 3.34 s\n",
      "Wall time: 3.34 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "nel = NEL()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "left_context, mention, right_context = (\n",
    "    \"American socialite\",\n",
    "    \"Paris\",\n",
    "    \"arrived in Turkish Cyprus on Aug. 4 for an appearance to promote her new perfume.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paris_Peace_Treaties,_1947 0.229919329286\n",
      "Paris 0.364550327261\n",
      "Paris_Peace_Conference 0.922771235307\n",
      "Paris–Nice 0.388005521148\n",
      "fr:Paris 0.198767289519\n",
      "Paris_Hilton 0.987708667914\n",
      "Catacombs_of_Paris 0.0474469698966\n",
      "Île-de-France 0.0764339289938\n",
      "Paris,_Ontario 0.193740231295\n",
      "Paris_syndrome 0.246216257413\n",
      "Panth%C3%A9on,_Paris 0.369571497043\n",
      "From_Paris_with_Love_(film) 0.606250623862\n",
      "Exposition_Internationale_des_Arts_et_Techniques_dans_la_Vie_Moderne 0.151203212639\n",
      "University_of_Paris 0.225954669217\n",
      "Notre_Dame_de_Paris 0.551950911681\n",
      "Paris_Peace_Accords 0.120936058462\n",
      "International_Exposition_(1867) 0.017266121693\n",
      "Fashion_week 0.597838769356\n",
      "Paris_(Malcolm_McLaren_album) 0.101084103187\n",
      "Paris_Saint-Germain_F.C. 0.542952239513\n",
      "Paris_sewers 0.0240961884459\n",
      "Paris_Combo 0.56066891551\n",
      "Dakar_Rally 0.0605802647769\n",
      "Biennale_de_Paris 0.0169372512028\n",
      "Paris–Brest 0.175204294423\n",
      "Paris_Sewer_Museum 0.1134365201\n",
      "15th_arrondissement_of_Paris 0.169871528943\n",
      "Fashion%20week 0.735459824403\n",
      "American_University_of_Paris 0.263087218006\n",
      "The_Paris_Review 0.450144976377\n",
      "Pavillon_de_Paris 0.0379617751266\n",
      "Dimitri_from_Paris 0.525669236978\n",
      "Jerry_Paris 0.56264296174\n",
      "Paris_(rapper) 0.613518814246\n",
      "9th_arrondissement_of_Paris 0.132725284745\n",
      "Conservatoire_de_Paris 0.0436758970221\n"
     ]
    }
   ],
   "source": [
    "nel.disabiguate(left_context, mention, right_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}