notebooks/try_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "7a806e00",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "4e0a3a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# APIs\n",
    "biencoder = 'http://localhost:30300/api/blink/biencoder' # mention # entity\n",
    "biencoder_mention = f'{biencoder}/mention'\n",
    "biencoder_entity = f'{biencoder}/entity'\n",
    "crossencoder = 'http://localhost:30302/api/blink/crossencoder'\n",
    "indexer = 'http://localhost:30301/api/indexer' # search # add\n",
    "indexer_search = f'{indexer}/search'\n",
    "indexer_add = f'{indexer}/add'\n",
    "indexer_reset = f'{indexer}/reset/rw'\n",
    "nilpredictor = 'http://localhost:30303/api/nilprediction'\n",
    "nilcluster = 'http://localhost:30305/api/nilcluster'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3308eb45",
   "metadata": {},
   "source": [
    "# Reset the New Kb (RW index) to ensure it is empty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "56cb5dab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Resetting RW index...\n",
      "Reset done.\n"
     ]
    }
   ],
   "source": [
    "print('Resetting RW index...')\n",
    "res_reset = requests.post(indexer_reset, data={})\n",
    "\n",
    "if res_reset.ok:\n",
    "    print('Reset done.')\n",
    "else:\n",
    "    print('ERROR while resetting!')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b74f80f1",
   "metadata": {},
   "source": [
    "Load some data (e.g. the first batch of the dev set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bd221dab",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.read_json('../incremental_dataset/dev/dev_0.jsonl', lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2e49b8db",
   "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>Wikipedia_ID</th>\n",
       "      <th>Wikipedia_title</th>\n",
       "      <th>Wikipedia_URL</th>\n",
       "      <th>left_context</th>\n",
       "      <th>mention_as_list</th>\n",
       "      <th>right_context</th>\n",
       "      <th>context_left</th>\n",
       "      <th>mention</th>\n",
       "      <th>context_right</th>\n",
       "      <th>y_category</th>\n",
       "      <th>...</th>\n",
       "      <th>docId</th>\n",
       "      <th>label_id</th>\n",
       "      <th>freebaseID</th>\n",
       "      <th>y_wikiurl_dump</th>\n",
       "      <th>query_id</th>\n",
       "      <th>p_formula</th>\n",
       "      <th>p_uniform</th>\n",
       "      <th>NIL</th>\n",
       "      <th>freq</th>\n",
       "      <th>batch</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3871014</td>\n",
       "      <td>Rainbow</td>\n",
       "      <td>http://en.wikipedia.org/wiki/Rainbow</td>\n",
       "      <td>[the, rainbow, '', from, the, wizard, of, oz, ...</td>\n",
       "      <td>[rainbow]</td>\n",
       "      <td>[,, a, symbol, of, good, ,, the, calm, after, ...</td>\n",
       "      <td>the Rainbow\" from the Wizard of Oz where Dorot...</td>\n",
       "      <td>Rainbow</td>\n",
       "      <td>, a symbol of good, the calm after the storm a...</td>\n",
       "      <td>[lucky symbols, atmospheric optical phenomena,...</td>\n",
       "      <td>...</td>\n",
       "      <td>5725443</td>\n",
       "      <td>http://hierroglyphic.blogspot.com/2010_09_01_a...</td>\n",
       "      <td>Some(9202a8c04000641f800000000038768c)</td>\n",
       "      <td>https://en.wikipedia.org/wiki?curid=3871014</td>\n",
       "      <td>4206740</td>\n",
       "      <td>3.162278e-95</td>\n",
       "      <td>0.734306</td>\n",
       "      <td>False</td>\n",
       "      <td>378</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>24997985</td>\n",
       "      <td>Tortoiseshell cat</td>\n",
       "      <td>http://en.wikipedia.org/wiki/Tortoiseshell_cat</td>\n",
       "      <td>[but, over, time, ,, this, characteristic, has...</td>\n",
       "      <td>[tortoiseshell]</td>\n",
       "      <td>[,, blue, ,, and, tabby, ., the, persian, is, ...</td>\n",
       "      <td>but over time, this characteristic has become ...</td>\n",
       "      <td>tortoiseshell</td>\n",
       "      <td>, blue, and tabby . The Persian is generally d...</td>\n",
       "      <td>[cat coat types]</td>\n",
       "      <td>...</td>\n",
       "      <td>8618624</td>\n",
       "      <td>http://www.sulit.com.ph/index.php/view+classif...</td>\n",
       "      <td>Some(9202a8c04000641f80000000000785ee)</td>\n",
       "      <td>https://en.wikipedia.org/wiki?curid=24997985</td>\n",
       "      <td>4227303</td>\n",
       "      <td>5.623400e-06</td>\n",
       "      <td>0.497273</td>\n",
       "      <td>False</td>\n",
       "      <td>21</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>43468</td>\n",
       "      <td>T. E. Lawrence</td>\n",
       "      <td>http://en.wikipedia.org/wiki/T._E._Lawrence</td>\n",
       "      <td>[some, of, the, political, implications, of, t...</td>\n",
       "      <td>[lawrence, of, arabia]</td>\n",
       "      <td>[,, was, a, badass, ., this, much, is, evident...</td>\n",
       "      <td>some of the political implications of the ways...</td>\n",
       "      <td>Lawrence of Arabia</td>\n",
       "      <td>, was a badass. This much is evident from any ...</td>\n",
       "      <td>[british army general list officers, british a...</td>\n",
       "      <td>...</td>\n",
       "      <td>6853281</td>\n",
       "      <td>http://joelsbookshelf.blogspot.com/2012/08/rea...</td>\n",
       "      <td>Some(9202a8c04000641f80000000000570ba)</td>\n",
       "      <td>https://en.wikipedia.org/wiki?curid=43468</td>\n",
       "      <td>4353903</td>\n",
       "      <td>1.000000e-02</td>\n",
       "      <td>0.633318</td>\n",
       "      <td>False</td>\n",
       "      <td>6</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>113519</td>\n",
       "      <td>Short (finance)</td>\n",
       "      <td>http://en.wikipedia.org/wiki/Short_(finance)</td>\n",
       "      <td>[of, homes, and, their, list-and-sold, price.a...</td>\n",
       "      <td>[short, sales]</td>\n",
       "      <td>[,, not, a, regular, sale, ), status, address,...</td>\n",
       "      <td>of homes and their list-and-sold price.Â  DOM ...</td>\n",
       "      <td>short sales</td>\n",
       "      <td>, not a regular sale) Status Address Bedroom, ...</td>\n",
       "      <td>[introductions, dutch inventions, 17th-century...</td>\n",
       "      <td>...</td>\n",
       "      <td>1557206</td>\n",
       "      <td>http://raymondong.com/</td>\n",
       "      <td>Some(9202a8c04000641f80000000000c73b4)</td>\n",
       "      <td>https://en.wikipedia.org/wiki?curid=113519</td>\n",
       "      <td>2138095</td>\n",
       "      <td>5.623413e-106</td>\n",
       "      <td>0.976293</td>\n",
       "      <td>False</td>\n",
       "      <td>421</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>156045</td>\n",
       "      <td>M*A*S*H (TV series)</td>\n",
       "      <td>http://en.wikipedia.org/wiki/M*A*S*H_(TV_series)</td>\n",
       "      <td>[watching, sporting, events, ,, season, finale...</td>\n",
       "      <td>[m*a*s*h]</td>\n",
       "      <td>[stood, as, the, most, watched, american, tv, ...</td>\n",
       "      <td>watching sporting events, season finales, and ...</td>\n",
       "      <td>M*A*S*H</td>\n",
       "      <td>stood as the most watched American TV broadcas...</td>\n",
       "      <td>[american, television programs based on films,...</td>\n",
       "      <td>...</td>\n",
       "      <td>10409920</td>\n",
       "      <td>http://ustelevision.com/2011/05/18/what-does-t...</td>\n",
       "      <td>Some(9202a8c04000641f8000000000123a15)</td>\n",
       "      <td>https://en.wikipedia.org/wiki?curid=156045</td>\n",
       "      <td>354180</td>\n",
       "      <td>3.162278e-01</td>\n",
       "      <td>0.760256</td>\n",
       "      <td>False</td>\n",
       "      <td>28</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 21 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   Wikipedia_ID      Wikipedia_title  \\\n",
       "0       3871014              Rainbow   \n",
       "1      24997985    Tortoiseshell cat   \n",
       "2         43468       T. E. Lawrence   \n",
       "3        113519      Short (finance)   \n",
       "4        156045  M*A*S*H (TV series)   \n",
       "\n",
       "                                      Wikipedia_URL  \\\n",
       "0              http://en.wikipedia.org/wiki/Rainbow   \n",
       "1    http://en.wikipedia.org/wiki/Tortoiseshell_cat   \n",
       "2       http://en.wikipedia.org/wiki/T._E._Lawrence   \n",
       "3      http://en.wikipedia.org/wiki/Short_(finance)   \n",
       "4  http://en.wikipedia.org/wiki/M*A*S*H_(TV_series)   \n",
       "\n",
       "                                        left_context         mention_as_list  \\\n",
       "0  [the, rainbow, '', from, the, wizard, of, oz, ...               [rainbow]   \n",
       "1  [but, over, time, ,, this, characteristic, has...         [tortoiseshell]   \n",
       "2  [some, of, the, political, implications, of, t...  [lawrence, of, arabia]   \n",
       "3  [of, homes, and, their, list-and-sold, price.a...          [short, sales]   \n",
       "4  [watching, sporting, events, ,, season, finale...               [m*a*s*h]   \n",
       "\n",
       "                                       right_context  \\\n",
       "0  [,, a, symbol, of, good, ,, the, calm, after, ...   \n",
       "1  [,, blue, ,, and, tabby, ., the, persian, is, ...   \n",
       "2  [,, was, a, badass, ., this, much, is, evident...   \n",
       "3  [,, not, a, regular, sale, ), status, address,...   \n",
       "4  [stood, as, the, most, watched, american, tv, ...   \n",
       "\n",
       "                                        context_left             mention  \\\n",
       "0  the Rainbow\" from the Wizard of Oz where Dorot...             Rainbow   \n",
       "1  but over time, this characteristic has become ...       tortoiseshell   \n",
       "2  some of the political implications of the ways...  Lawrence of Arabia   \n",
       "3  of homes and their list-and-sold price.Â  DOM ...         short sales   \n",
       "4  watching sporting events, season finales, and ...             M*A*S*H   \n",
       "\n",
       "                                       context_right  \\\n",
       "0  , a symbol of good, the calm after the storm a...   \n",
       "1  , blue, and tabby . The Persian is generally d...   \n",
       "2  , was a badass. This much is evident from any ...   \n",
       "3  , not a regular sale) Status Address Bedroom, ...   \n",
       "4  stood as the most watched American TV broadcas...   \n",
       "\n",
       "                                          y_category  ...     docId  \\\n",
       "0  [lucky symbols, atmospheric optical phenomena,...  ...   5725443   \n",
       "1                                   [cat coat types]  ...   8618624   \n",
       "2  [british army general list officers, british a...  ...   6853281   \n",
       "3  [introductions, dutch inventions, 17th-century...  ...   1557206   \n",
       "4  [american, television programs based on films,...  ...  10409920   \n",
       "\n",
       "                                            label_id  \\\n",
       "0  http://hierroglyphic.blogspot.com/2010_09_01_a...   \n",
       "1  http://www.sulit.com.ph/index.php/view+classif...   \n",
       "2  http://joelsbookshelf.blogspot.com/2012/08/rea...   \n",
       "3                             http://raymondong.com/   \n",
       "4  http://ustelevision.com/2011/05/18/what-does-t...   \n",
       "\n",
       "                               freebaseID  \\\n",
       "0  Some(9202a8c04000641f800000000038768c)   \n",
       "1  Some(9202a8c04000641f80000000000785ee)   \n",
       "2  Some(9202a8c04000641f80000000000570ba)   \n",
       "3  Some(9202a8c04000641f80000000000c73b4)   \n",
       "4  Some(9202a8c04000641f8000000000123a15)   \n",
       "\n",
       "                                 y_wikiurl_dump query_id      p_formula  \\\n",
       "0   https://en.wikipedia.org/wiki?curid=3871014  4206740   3.162278e-95   \n",
       "1  https://en.wikipedia.org/wiki?curid=24997985  4227303   5.623400e-06   \n",
       "2     https://en.wikipedia.org/wiki?curid=43468  4353903   1.000000e-02   \n",
       "3    https://en.wikipedia.org/wiki?curid=113519  2138095  5.623413e-106   \n",
       "4    https://en.wikipedia.org/wiki?curid=156045   354180   3.162278e-01   \n",
       "\n",
       "   p_uniform    NIL  freq  batch  \n",
       "0   0.734306  False   378      0  \n",
       "1   0.497273  False    21      0  \n",
       "2   0.633318  False     6      0  \n",
       "3   0.976293  False   421      0  \n",
       "4   0.760256  False    28      0  \n",
       "\n",
       "[5 rows x 21 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4cc54ba",
   "metadata": {},
   "source": [
    "Select 100 not-NIL + 100 NIL mentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "4b84f295",
   "metadata": {},
   "outputs": [],
   "source": [
    "number_not_nil = 50\n",
    "number_nil = 50\n",
    "\n",
    "selection = data.query('~NIL').head(number_not_nil)\n",
    "selection = pd.concat([selection, data.query('NIL').head(number_nil)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "a6987569",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 21)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "selection.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e13acde0",
   "metadata": {},
   "source": [
    "Check how many NIL mentions should be clustered together (since they refer to the same out-of-KB entity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "14e90bd0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "262831      5\n",
       "4607980     4\n",
       "21980       3\n",
       "182494      3\n",
       "145699      2\n",
       "14900       2\n",
       "43999623    2\n",
       "30310       1\n",
       "5869719     1\n",
       "22316162    1\n",
       "1926443     1\n",
       "4565664     1\n",
       "15389730    1\n",
       "14508500    1\n",
       "94834       1\n",
       "7076247     1\n",
       "425554      1\n",
       "207750      1\n",
       "23281       1\n",
       "40359       1\n",
       "262891      1\n",
       "501582      1\n",
       "19317904    1\n",
       "978740      1\n",
       "755645      1\n",
       "43452       1\n",
       "3185609     1\n",
       "28387289    1\n",
       "12639363    1\n",
       "27705200    1\n",
       "668081      1\n",
       "2135896     1\n",
       "45492       1\n",
       "363196      1\n",
       "1619786     1\n",
       "4223084     1\n",
       "Name: Wikipedia_ID, dtype: int64"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "selection.query('NIL')['Wikipedia_ID'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "ab2a092e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 21)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "selection.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30e11499",
   "metadata": {},
   "source": [
    "# Biencoder encode mentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "15cc3c59",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_biencoder = requests.post(biencoder_mention,\n",
    "        json=selection[[\n",
    "            'mention',\n",
    "            'context_left',\n",
    "            'context_right'\n",
    "            ]].to_dict(orient='records'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "759b888d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Biencode OK\n",
      "Encoded 100 entities.\n"
     ]
    }
   ],
   "source": [
    "if res_biencoder.ok:\n",
    "    selection['encoding'] = res_biencoder.json()['encodings']\n",
    "    print('Biencode OK')\n",
    "    print('Encoded {} entities.'.format(selection.shape[0]))\n",
    "else:\n",
    "    print('Biencoder ERROR')\n",
    "    print(res_biencoder)\n",
    "    raise Exception('Biencoder ERROR')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b235bf",
   "metadata": {},
   "source": [
    "The encoding columns has been added with the base64 encoded vector representing the mention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "539d8db3",
   "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>encoding</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Y1J+PkX+ADw1DLO8hd7evFtK5L1xXSA9AlKgPCnUwTyp7D...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>+4FqPeIoGj6bioQ+3jnJvVQpfr1r7qq96fFQPZ5IvD0dcR...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>QhF3Po2lADyk5gs+11sVvXNHDb0Ts0c8oUe5vGqJwr1u+u...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>MACiPZd4aD4suke8wZMfPEs5zb1YiBS+SZ5ZvOg2wD0mPW...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>TUqpPh9cET4XdAQ+pVUlPl3rGT7PtSU+PmdpvswESTzFMQ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            encoding\n",
       "0  Y1J+PkX+ADw1DLO8hd7evFtK5L1xXSA9AlKgPCnUwTyp7D...\n",
       "1  +4FqPeIoGj6bioQ+3jnJvVQpfr1r7qq96fFQPZ5IvD0dcR...\n",
       "2  QhF3Po2lADyk5gs+11sVvXNHDb0Ts0c8oUe5vGqJwr1u+u...\n",
       "3  MACiPZd4aD4suke8wZMfPEs5zb1YiBS+SZ5ZvOg2wD0mPW...\n",
       "4  TUqpPh9cET4XdAQ+pVUlPl3rGT7PtSU+PmdpvswESTzFMQ..."
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "selection[['encoding']].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06ae6a4f",
   "metadata": {},
   "source": [
    "# Retrieval with indexer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "65309931",
   "metadata": {},
   "outputs": [],
   "source": [
    "body = {\n",
    "    'encodings': selection['encoding'].values.tolist(),\n",
    "    'top_k': 10 # top_10 candidates\n",
    "}\n",
    "res_indexer = requests.post(indexer_search, json=body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "7bf1e680",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexer OK\n"
     ]
    }
   ],
   "source": [
    "if res_indexer.ok:\n",
    "    candidates = res_indexer.json()\n",
    "    print('Indexer OK')\n",
    "else:\n",
    "    print('ERROR with the indexer.')\n",
    "    print(res_indexer)\n",
    "    print(res_indexer.json())\n",
    "\n",
    "if len(candidates) == 0 or len(candidates[0]) == 0:\n",
    "    print('No candidates received.')\n",
    "\n",
    "selection['candidates'] = candidates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "c4290d9d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'raw_score': 322.6261901855469,\n",
       "  'id': 849365,\n",
       "  'wikipedia_id': 3871014,\n",
       "  'title': 'Rainbow',\n",
       "  'url': 'https://en.wikipedia.org/wiki?curid=3871014',\n",
       "  'type_': None,\n",
       "  'indexer': 10,\n",
       "  'score': 83.38381958007812,\n",
       "  'norm_score': 0.5085703801599667},\n",
       " {'raw_score': 327.61798095703125,\n",
       "  'id': 331878,\n",
       "  'wikipedia_id': 1026203,\n",
       "  'title': 'Rainbows in mythology',\n",
       "  'url': 'https://en.wikipedia.org/wiki?curid=1026203',\n",
       "  'type_': None,\n",
       "  'indexer': 10,\n",
       "  'score': 80.88790893554688,\n",
       "  'norm_score': 0.5350073800335146},\n",
       " {'raw_score': 330.5491943359375,\n",
       "  'id': 35028,\n",
       "  'wikipedia_id': 71079,\n",
       "  'title': 'Rainbow flag',\n",
       "  'url': 'https://en.wikipedia.org/wiki?curid=71079',\n",
       "  'type_': None,\n",
       "  'indexer': 10,\n",
       "  'score': 79.42232513427734,\n",
       "  'norm_score': 0.4614385006898517}]"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 3 candidates\n",
    "selection['candidates'].iloc[0][:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "338a6ac9",
   "metadata": {},
   "source": [
    "# NIL prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "3091720b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_for_nil_prediction(x, mention='mention'):\n",
    "    \"\"\"\n",
    "    Function to prepare the features required by the nil predictor\n",
    "    \"\"\"\n",
    "    c = x['candidates']\n",
    "\n",
    "    is_nil = False\n",
    "    features = {}\n",
    "\n",
    "    if len(c) == 0:\n",
    "        is_nil = True\n",
    "        return is_nil, features\n",
    "\n",
    "    is_cross = 'is_cross' in c[0] and c[0]['is_cross']\n",
    "\n",
    "    features = {}\n",
    "    if not is_cross:\n",
    "        # bi only\n",
    "        features['max_bi'] = c[0]['score']\n",
    "    else:\n",
    "        # cross\n",
    "        if 'bi_score' in c[0]:\n",
    "            features['max_bi'] = c[0]['bi_score']\n",
    "        features['max_cross'] = c[0]['score']\n",
    "\n",
    "    features['mention'] = x[mention]\n",
    "    features['title'] = c[0]['title']\n",
    "    features['topcandidates'] = c\n",
    "\n",
    "    return is_nil, features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "35484598",
   "metadata": {},
   "outputs": [],
   "source": [
    "selection[['is_nil', 'nil_features']] = selection.apply(prepare_for_nil_prediction, axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "beec6245",
   "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>is_nil</th>\n",
       "      <th>nil_features</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>False</td>\n",
       "      <td>{'max_bi': 83.38381958007812, 'mention': 'Rain...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>False</td>\n",
       "      <td>{'max_bi': 84.1929931640625, 'mention': 'torto...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>False</td>\n",
       "      <td>{'max_bi': 78.03597259521484, 'mention': 'Lawr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>False</td>\n",
       "      <td>{'max_bi': 83.60071563720703, 'mention': 'shor...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>False</td>\n",
       "      <td>{'max_bi': 80.76179504394531, 'mention': 'M*A*...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   is_nil                                       nil_features\n",
       "0   False  {'max_bi': 83.38381958007812, 'mention': 'Rain...\n",
       "1   False  {'max_bi': 84.1929931640625, 'mention': 'torto...\n",
       "2   False  {'max_bi': 78.03597259521484, 'mention': 'Lawr...\n",
       "3   False  {'max_bi': 83.60071563720703, 'mention': 'shor...\n",
       "4   False  {'max_bi': 80.76179504394531, 'mention': 'M*A*..."
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "selection[['is_nil', 'nil_features']].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "e933e910",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NIL pred OK\n",
      "Estimated 86 entities as NOT NIL\n",
      "Estimated 14 entities as NIL\n"
     ]
    }
   ],
   "source": [
    "## NIL prediction\n",
    "# initialize fields (default NIL)\n",
    "selection['nil_score'] = np.zeros(selection.shape[0])\n",
    "not_yet_nil = selection.query('is_nil == False')\n",
    "\n",
    "if not_yet_nil.shape[0] > 0:\n",
    "    res_nilpredictor = requests.post(nilpredictor, json=not_yet_nil['nil_features'].values.tolist())\n",
    "    if res_nilpredictor.ok:\n",
    "        print('NIL pred OK')\n",
    "        nil_scores_bi = np.array(res_nilpredictor.json()['nil_score_bi'])\n",
    "    else:\n",
    "        print('ERROR during NIL prediction')\n",
    "        print(res_nilpredictor)\n",
    "        print(res_nilpredictor.json())\n",
    "else:\n",
    "    print('ERROR. Probably the KB is emtpy')\n",
    "\n",
    "selection.loc[not_yet_nil.index, 'nil_score'] = nil_scores_bi\n",
    "\n",
    "nil_threshold = 0.5\n",
    "# if below threshold --> is NIL\n",
    "selection['is_nil'] = selection['nil_score'].apply(lambda x: x < nil_threshold)\n",
    "\n",
    "print('Estimated {} entities as NOT NIL'.format(selection.eval('is_nil == False').sum()))\n",
    "print('Estimated {} entities as NIL'.format(selection.eval('is_nil == True').sum()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "609109f7",
   "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>Wikipedia_title</th>\n",
       "      <th>mention</th>\n",
       "      <th>NIL</th>\n",
       "      <th>is_nil</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Rainbow</td>\n",
       "      <td>Rainbow</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Tortoiseshell cat</td>\n",
       "      <td>tortoiseshell</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Short (finance)</td>\n",
       "      <td>short sales</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>M*A*S*H (TV series)</td>\n",
       "      <td>M*A*S*H</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Limbo</td>\n",
       "      <td>limbo</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       Wikipedia_title        mention    NIL  is_nil\n",
       "0              Rainbow        Rainbow  False   False\n",
       "1    Tortoiseshell cat  tortoiseshell  False   False\n",
       "3      Short (finance)    short sales  False   False\n",
       "4  M*A*S*H (TV series)        M*A*S*H  False   False\n",
       "5                Limbo          limbo  False   False"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NIL\n",
    "selection.query('~is_nil')[[\"Wikipedia_title\", \"mention\", \"NIL\", \"is_nil\"]].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "0e8d35e2",
   "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>Wikipedia_title</th>\n",
       "      <th>mention</th>\n",
       "      <th>NIL</th>\n",
       "      <th>is_nil</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>T. E. Lawrence</td>\n",
       "      <td>Lawrence of Arabia</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>Thebes, Greece</td>\n",
       "      <td>Thebes</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>DSM-IV codes</td>\n",
       "      <td>DSM-IV</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>Myra Breckinridge (film)</td>\n",
       "      <td>Myra Breckinridge</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1183</th>\n",
       "      <td>Arcadia Group</td>\n",
       "      <td>arcadia</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               Wikipedia_title             mention    NIL  is_nil\n",
       "2               T. E. Lawrence  Lawrence of Arabia  False    True\n",
       "37              Thebes, Greece              Thebes  False    True\n",
       "41                DSM-IV codes              DSM-IV  False    True\n",
       "43    Myra Breckinridge (film)   Myra Breckinridge  False    True\n",
       "1183             Arcadia Group             arcadia   True    True"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NIL\n",
    "selection.query('is_nil')[[\"Wikipedia_title\", \"mention\", \"NIL\", \"is_nil\"]].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c5c6982",
   "metadata": {},
   "source": [
    "# NIL Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "7e84725f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NIL cluster OK\n"
     ]
    }
   ],
   "source": [
    "## Entity Clustering\n",
    "nil_mentions = selection.query('is_nil == True')\n",
    "\n",
    "res_nilcluster = requests.post(nilcluster, json={\n",
    "        'ids': nil_mentions.index.tolist(),\n",
    "        'mentions': nil_mentions[\"mention\"].values.tolist(),\n",
    "        'encodings': nil_mentions['encoding'].values.tolist()\n",
    "    })\n",
    "\n",
    "if not res_nilcluster.ok:\n",
    "    print('NIL cluster ERROR')\n",
    "else:\n",
    "    print('NIL cluster OK')\n",
    "\n",
    "clusters = pd.DataFrame(res_nilcluster.json())\n",
    "\n",
    "# visualize big clusters first \n",
    "clusters = clusters.sort_values(by='nelements', ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "8c3a6184",
   "metadata": {
    "scrolled": true
   },
   "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>title</th>\n",
       "      <th>nelements</th>\n",
       "      <th>mentions_id</th>\n",
       "      <th>mentions</th>\n",
       "      <th>center</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Miami Vice</td>\n",
       "      <td>5</td>\n",
       "      <td>[3842, 4281, 4708, 5163, 8710]</td>\n",
       "      <td>[Miami Vice, Miami Vice, Miami Vice, Miami Vic...</td>\n",
       "      <td>69xJPhww1jsoZmy9kIqSPgLgvbxekgk+gGaLvkxoFL1Vv9...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Lawrence of Arabia</td>\n",
       "      <td>2</td>\n",
       "      <td>[2, 4715]</td>\n",
       "      <td>[Lawrence of Arabia, Lawrence of Arabia]</td>\n",
       "      <td>QhF3Po2lADyk5gs+11sVvXNHDb0Ts0c8oUe5vGqJwr1u+u...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>DSM-IV</td>\n",
       "      <td>1</td>\n",
       "      <td>[41]</td>\n",
       "      <td>[DSM-IV]</td>\n",
       "      <td>jTKIPtQyOT0ViYw9b2oMvtA03r1xCeG8JoorvlaS+b3eGU...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Thebes</td>\n",
       "      <td>1</td>\n",
       "      <td>[37]</td>\n",
       "      <td>[Thebes]</td>\n",
       "      <td>lfWtPdLcdD1B2wQ+kh+Qu0DcpL0GjzI9HR4hvgldzb1zvF...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>arcadia</td>\n",
       "      <td>1</td>\n",
       "      <td>[1183]</td>\n",
       "      <td>[arcadia]</td>\n",
       "      <td>EXcdPjrcYj360be9WuifvUk0T72bp8W98yKFvI/KMT0H7c...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                title  nelements                     mentions_id  \\\n",
       "7          Miami Vice          5  [3842, 4281, 4708, 5163, 8710]   \n",
       "5  Lawrence of Arabia          2                       [2, 4715]   \n",
       "0              DSM-IV          1                            [41]   \n",
       "1              Thebes          1                            [37]   \n",
       "2             arcadia          1                          [1183]   \n",
       "\n",
       "                                            mentions  \\\n",
       "7  [Miami Vice, Miami Vice, Miami Vice, Miami Vic...   \n",
       "5           [Lawrence of Arabia, Lawrence of Arabia]   \n",
       "0                                           [DSM-IV]   \n",
       "1                                           [Thebes]   \n",
       "2                                          [arcadia]   \n",
       "\n",
       "                                              center  \n",
       "7  69xJPhww1jsoZmy9kIqSPgLgvbxekgk+gGaLvkxoFL1Vv9...  \n",
       "5  QhF3Po2lADyk5gs+11sVvXNHDb0Ts0c8oUe5vGqJwr1u+u...  \n",
       "0  jTKIPtQyOT0ViYw9b2oMvtA03r1xCeG8JoorvlaS+b3eGU...  \n",
       "1  lfWtPdLcdD1B2wQ+kh+Qu0DcpL0GjzI9HR4hvgldzb1zvF...  \n",
       "2  EXcdPjrcYj360be9WuifvUk0T72bp8W98yKFvI/KMT0H7c...  "
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clusters.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e40940bc",
   "metadata": {},
   "source": [
    "# Add New entities to the New KB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "d9582d59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new entities added correctly.\n"
     ]
    }
   ],
   "source": [
    "selection_new = clusters[['title', 'center']].rename(columns={'center': 'encoding', 'mode': 'wikipedia_id'})\n",
    "new_indexed = requests.post(indexer_add, json=selection_new.to_dict(orient='records'))\n",
    "\n",
    "if not new_indexed.ok:\n",
    "    print('error adding new entities')\n",
    "else:\n",
    "    print('new entities added correctly.')\n",
    "    new_indexed = new_indexed.json()\n",
    "    clusters['index_id'] = new_indexed['ids']\n",
    "    clusters['index_indexer'] = new_indexed['indexer']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9396496",
   "metadata": {},
   "source": [
    "At this point these new entities are retrieved in the indexer step so that subsequent documents that mentions them could be linked."
   ]
  }
 ],
 "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": 5
}