GraSCCo_Evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "079d2b75-7bd7-4cee-a8f3-c4d59bb52f7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from xmen.linkers import default_ensemble\n",
    "from xmen.data import filter_and_apply_threshold\n",
    "from xmen.evaluation import evaluate, evaluate_at_k, error_analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e0c5fe2-5302-40d9-b9e1-c3a911dac18d",
   "metadata": {},
   "source": [
    "## Load Dataset in BigBIO Format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4da15404-5c2a-4ffd-829c-c9bbc606c33f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fatal: destination path 'biomedical' already exists and is not an empty directory.\n"
     ]
    }
   ],
   "source": [
    "# Clone forked repo until loader with SNOMED CT layer is on HF Hub\n",
    "!git clone -b grascco_snomed git@github.com:phlobo/biomedical.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a25a9eb2-e86c-40c3-b23e-b3008a3d19fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "ground_truth = datasets.load_dataset(\n",
    "    'biomedical/bigbio/hub/hub_repos/grascco/grascco.py', \n",
    "    'grascco_snomed_bigbio_kb',\n",
    "    data_dir='../gemtex_oncology/annotation/json/'\n",
    ")['train']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67397d42-3c8d-4154-9c0a-64c7ed5d2920",
   "metadata": {},
   "source": [
    "## Generate Candidates Using xMEN Ensemble Linker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fa6290fb-5010-4621-b071-23b1f2d76da6",
   "metadata": {},
   "outputs": [
    {
     "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\">[11/06/24 10:52:59] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loading hierarchical faiss index                                <a href=\"file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/sap_bert_linker.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">sap_bert_linker.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/sap_bert_linker.py#153\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">153</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/06/24 10:52:59]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading hierarchical faiss index                                \u001b]8;id=404306;file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/sap_bert_linker.py\u001b\\\u001b[2msap_bert_linker.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=592440;file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/sap_bert_linker.py#153\u001b\\\u001b[2m153\u001b[0m\u001b]8;;\u001b\\\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 index from                                                 <a href=\"file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">faiss_indexer.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py#64\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">64</span></a>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         xmen_index/index/sapbert/embed_faiss_hier.pickle                   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading index from                                                 \u001b]8;id=689352;file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=525580;file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py#64\u001b\\\u001b[2m64\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         xmen_index/index/sapbert/embed_faiss_hier.pickle                   \u001b[2m                   \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\">[11/06/24 10:53:03] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loaded index of type <span style=\"font-weight: bold\">&lt;</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">class</span><span style=\"color: #000000; text-decoration-color: #000000\"> </span><span style=\"color: #008000; text-decoration-color: #008000\">'faiss.swigfaiss.IndexHNSWFlat'</span><span style=\"font-weight: bold\">&gt;</span> and   <a href=\"file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">faiss_indexer.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py#66\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">66</span></a>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         size <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1967771</span>                                                       <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/06/24 10:53:03]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loaded index of type \u001b[1m<\u001b[0m\u001b[1;95mclass\u001b[0m\u001b[39m \u001b[0m\u001b[32m'faiss.swigfaiss.IndexHNSWFlat'\u001b[0m\u001b[1m>\u001b[0m and   \u001b]8;id=341769;file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=999351;file:///Users/phlobo/miniconda3/envs/xmen/lib/python3.10/site-packages/xmen/linkers/faiss_indexer.py#66\u001b\\\u001b[2m66\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         size \u001b[1;36m1967771\u001b[0m                                                       \u001b[2m                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "linker = default_ensemble('xmen_index/index/', cuda=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1cc69f99-a209-474a-b905-bf6619991a76",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:05<00:00,  5.22s/ examples]\n"
     ]
    }
   ],
   "source": [
    "from utils import handle_dates, to_map_fn\n",
    "\n",
    "candidates_ = linker.predict_batch(ground_truth)\n",
    "candidates = candidates_.map(to_map_fn(handle_dates))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3def9cc-52f7-463a-a70d-40ef1c4aa7c1",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e64fe65b-dd4e-466c-8cb4-bf11a10f063f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.546875\n",
      "Recall@2 0.609375\n",
      "Recall@3 0.65625\n",
      "Recall@4 0.671875\n"
     ]
    }
   ],
   "source": [
    "_ = evaluate_at_k(ground_truth, candidates, eval_k=[1,2,3,4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "17ec8730-ce95-42cf-990b-2c14ac3b8576",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'strict': {'precision': 0.546875,\n",
       "  'recall': 0.546875,\n",
       "  'fscore': 0.546875,\n",
       "  'ptp': 35,\n",
       "  'fp': 29,\n",
       "  'rtp': 35,\n",
       "  'fn': 29,\n",
       "  'n_docs_system': 1,\n",
       "  'n_annos_system': 64,\n",
       "  'n_docs_gold': 1,\n",
       "  'n_annos_gold': 64}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluate(ground_truth, candidates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "94148620-661c-4353-b45a-3a751276aa52",
   "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>_word_len</th>\n",
       "      <th>_abbrev</th>\n",
       "      <th>gt_start</th>\n",
       "      <th>gt_end</th>\n",
       "      <th>gt_text</th>\n",
       "      <th>gold_type</th>\n",
       "      <th>gold_concept</th>\n",
       "      <th>pred_index</th>\n",
       "      <th>pred_index_score</th>\n",
       "      <th>pred_top</th>\n",
       "      <th>pred_top_score</th>\n",
       "      <th>document_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>105</td>\n",
       "      <td>113</td>\n",
       "      <td>[4.4.1997]</td>\n",
       "      <td>Literal</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258695005'}</td>\n",
       "      <td>0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>129</td>\n",
       "      <td>134</td>\n",
       "      <td>[19.3.]</td>\n",
       "      <td>Literal</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258695005'}</td>\n",
       "      <td>0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>143</td>\n",
       "      <td>151</td>\n",
       "      <td>[7.5.2029]</td>\n",
       "      <td>Literal</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258695005'}</td>\n",
       "      <td>0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>215</td>\n",
       "      <td>226</td>\n",
       "      <td>[Verbrennung]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '125666000'}</td>\n",
       "      <td>2</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>48333001</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>227</td>\n",
       "      <td>229</td>\n",
       "      <td>[1.]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258351006'}</td>\n",
       "      <td>22</td>\n",
       "      <td>0.768320</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>232</td>\n",
       "      <td>241</td>\n",
       "      <td>[3. Grades]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258353009'}</td>\n",
       "      <td>1</td>\n",
       "      <td>0.934421</td>\n",
       "      <td>61026006</td>\n",
       "      <td>0.934421</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>243</td>\n",
       "      <td>247</td>\n",
       "      <td>[Kopf]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '302548004'}</td>\n",
       "      <td>5</td>\n",
       "      <td>0.806592</td>\n",
       "      <td>69536005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>250</td>\n",
       "      <td>254</td>\n",
       "      <td>[Hals]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '49928004'}</td>\n",
       "      <td>37</td>\n",
       "      <td>0.699021</td>\n",
       "      <td>45048000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>255</td>\n",
       "      <td>257</td>\n",
       "      <td>[5%]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '113341005'}</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1</td>\n",
       "      <td>True</td>\n",
       "      <td>260</td>\n",
       "      <td>263</td>\n",
       "      <td>[KOF]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '301898006'}</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>34763001</td>\n",
       "      <td>0.831992</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   _word_len  _abbrev  gt_start  gt_end        gt_text gold_type  \\\n",
       "0          1    False       105     113     [4.4.1997]   Literal   \n",
       "1          1    False       129     134        [19.3.]   Literal   \n",
       "2          1    False       143     151     [7.5.2029]   Literal   \n",
       "3          1    False       215     226  [Verbrennung]   Concept   \n",
       "4          1    False       227     229           [1.]   Concept   \n",
       "5          2    False       232     241    [3. Grades]   Concept   \n",
       "6          1    False       243     247         [Kopf]   Concept   \n",
       "7          1    False       250     254         [Hals]   Concept   \n",
       "8          1    False       255     257           [5%]   Concept   \n",
       "9          1     True       260     263          [KOF]   Concept   \n",
       "\n",
       "                                     gold_concept  pred_index  \\\n",
       "0  {'db_name': 'SNOMED CT', 'db_id': '258695005'}           0   \n",
       "1  {'db_name': 'SNOMED CT', 'db_id': '258695005'}           0   \n",
       "2  {'db_name': 'SNOMED CT', 'db_id': '258695005'}           0   \n",
       "3  {'db_name': 'SNOMED CT', 'db_id': '125666000'}           2   \n",
       "4  {'db_name': 'SNOMED CT', 'db_id': '258351006'}          22   \n",
       "5  {'db_name': 'SNOMED CT', 'db_id': '258353009'}           1   \n",
       "6  {'db_name': 'SNOMED CT', 'db_id': '302548004'}           5   \n",
       "7   {'db_name': 'SNOMED CT', 'db_id': '49928004'}          37   \n",
       "8  {'db_name': 'SNOMED CT', 'db_id': '113341005'}          -1   \n",
       "9  {'db_name': 'SNOMED CT', 'db_id': '301898006'}          -1   \n",
       "\n",
       "   pred_index_score   pred_top  pred_top_score document_id  \n",
       "0          1.000000  258695005        1.000000  Albers.txt  \n",
       "1          1.000000  258695005        1.000000  Albers.txt  \n",
       "2          1.000000  258695005        1.000000  Albers.txt  \n",
       "3          1.000000   48333001        1.000000  Albers.txt  \n",
       "4          0.768320  258695005        1.000000  Albers.txt  \n",
       "5          0.934421   61026006        0.934421  Albers.txt  \n",
       "6          0.806592   69536005        1.000000  Albers.txt  \n",
       "7          0.699021   45048000        1.000000  Albers.txt  \n",
       "8               NaN  258695005        1.000000  Albers.txt  \n",
       "9               NaN   34763001        0.831992  Albers.txt  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ea_df = error_analysis(ground_truth, candidates, tasks=['nen'])\n",
    "ea_df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e415a18-5200-44de-a76e-2b7b68eb61db",
   "metadata": {},
   "source": [
    "# Pre-trained Re-ranker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "67a722ef-88f5-4f59-baba-33de4b5796fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.reranking import CrossEncoderReranker\n",
    "from xmen import load_kb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ee8ef724-3b52-4fed-801a-f60949a21dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_candidates = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e643afaf-8e6d-4a0e-8989-7c40bef31345",
   "metadata": {},
   "outputs": [],
   "source": [
    "kb = load_kb('xmen_index/snomed_german.jsonl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8401cf5a-cf7f-45d3-a635-c7791fd26893",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context length: 128\n",
      "Use NIL values: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  4.97 examples/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████| 64/64 [00:00<00:00, 112693.31it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████| 64/64 [00:00<00:00, 592572.75it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 64/64 [00:00<00:00, 9419.78it/s]\n"
     ]
    }
   ],
   "source": [
    "ce_dataset = CrossEncoderReranker.prepare_data(candidates, ground_truth, kb, k=n_candidates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3b8bab73-fe1c-44ce-b89e-456e11d2b4f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr = CrossEncoderReranker.load('phlobo/xmen-de-ce-medmentions', device='mps') # works on MacOS M1, you may set device to cpu or <GPU ID>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0faa72d6-0c68-466d-a841-d10b5d1d220a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Batches: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 64/64 [00:19<00:00,  3.35it/s]\n",
      "Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  6.82 examples/s]\n",
      "Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 81.85 examples/s]\n",
      "Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 96.41 examples/s]\n"
     ]
    }
   ],
   "source": [
    "reranked_ = rr.rerank_batch(candidates, ce_dataset, k=n_candidates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "cac0a6fb-66ec-4b33-8e24-0798b95e0744",
   "metadata": {},
   "outputs": [],
   "source": [
    "reranked = reranked_.map(to_map_fn(handle_dates))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "612c5131-b257-4e67-93d7-4f6604bb48f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.609375\n",
      "Recall@2 0.640625\n",
      "Recall@3 0.65625\n",
      "Recall@4 0.671875\n"
     ]
    }
   ],
   "source": [
    "_ = evaluate_at_k(ground_truth, reranked, eval_k=[1,2,3,4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e2847978-2957-4f9b-9be4-52600cec0a3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'strict': {'precision': 0.6842105263157895,\n",
       "  'recall': 0.609375,\n",
       "  'fscore': 0.6446280991735538,\n",
       "  'ptp': 39,\n",
       "  'fp': 18,\n",
       "  'rtp': 39,\n",
       "  'fn': 25,\n",
       "  'n_docs_system': 1,\n",
       "  'n_annos_system': 57,\n",
       "  'n_docs_gold': 1,\n",
       "  'n_annos_gold': 64}}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluate(ground_truth, reranked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d719ded2-d54f-4d64-a285-d20748d1d90e",
   "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>_word_len</th>\n",
       "      <th>_abbrev</th>\n",
       "      <th>gt_start</th>\n",
       "      <th>gt_end</th>\n",
       "      <th>gt_text</th>\n",
       "      <th>gold_type</th>\n",
       "      <th>gold_concept</th>\n",
       "      <th>pred_index</th>\n",
       "      <th>pred_index_score</th>\n",
       "      <th>pred_top</th>\n",
       "      <th>pred_top_score</th>\n",
       "      <th>document_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>105</td>\n",
       "      <td>113</td>\n",
       "      <td>[4.4.1997]</td>\n",
       "      <td>Literal</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258695005'}</td>\n",
       "      <td>0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>129</td>\n",
       "      <td>134</td>\n",
       "      <td>[19.3.]</td>\n",
       "      <td>Literal</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258695005'}</td>\n",
       "      <td>0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>143</td>\n",
       "      <td>151</td>\n",
       "      <td>[7.5.2029]</td>\n",
       "      <td>Literal</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258695005'}</td>\n",
       "      <td>0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>215</td>\n",
       "      <td>226</td>\n",
       "      <td>[Verbrennung]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '125666000'}</td>\n",
       "      <td>1</td>\n",
       "      <td>0.166484</td>\n",
       "      <td>48333001</td>\n",
       "      <td>0.188725</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>227</td>\n",
       "      <td>229</td>\n",
       "      <td>[1.]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258351006'}</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>232</td>\n",
       "      <td>241</td>\n",
       "      <td>[3. Grades]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '258353009'}</td>\n",
       "      <td>0</td>\n",
       "      <td>0.141101</td>\n",
       "      <td>258353009</td>\n",
       "      <td>0.141101</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>243</td>\n",
       "      <td>247</td>\n",
       "      <td>[Kopf]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '302548004'}</td>\n",
       "      <td>5</td>\n",
       "      <td>0.114893</td>\n",
       "      <td>69536005</td>\n",
       "      <td>0.223935</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>250</td>\n",
       "      <td>254</td>\n",
       "      <td>[Hals]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '49928004'}</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>45048000</td>\n",
       "      <td>0.241662</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>255</td>\n",
       "      <td>257</td>\n",
       "      <td>[5%]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '113341005'}</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>258695005</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1</td>\n",
       "      <td>True</td>\n",
       "      <td>260</td>\n",
       "      <td>263</td>\n",
       "      <td>[KOF]</td>\n",
       "      <td>Concept</td>\n",
       "      <td>{'db_name': 'SNOMED CT', 'db_id': '301898006'}</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>34763001</td>\n",
       "      <td>0.138157</td>\n",
       "      <td>Albers.txt</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   _word_len  _abbrev  gt_start  gt_end        gt_text gold_type  \\\n",
       "0          1    False       105     113     [4.4.1997]   Literal   \n",
       "1          1    False       129     134        [19.3.]   Literal   \n",
       "2          1    False       143     151     [7.5.2029]   Literal   \n",
       "3          1    False       215     226  [Verbrennung]   Concept   \n",
       "4          1    False       227     229           [1.]   Concept   \n",
       "5          2    False       232     241    [3. Grades]   Concept   \n",
       "6          1    False       243     247         [Kopf]   Concept   \n",
       "7          1    False       250     254         [Hals]   Concept   \n",
       "8          1    False       255     257           [5%]   Concept   \n",
       "9          1     True       260     263          [KOF]   Concept   \n",
       "\n",
       "                                     gold_concept  pred_index  \\\n",
       "0  {'db_name': 'SNOMED CT', 'db_id': '258695005'}           0   \n",
       "1  {'db_name': 'SNOMED CT', 'db_id': '258695005'}           0   \n",
       "2  {'db_name': 'SNOMED CT', 'db_id': '258695005'}           0   \n",
       "3  {'db_name': 'SNOMED CT', 'db_id': '125666000'}           1   \n",
       "4  {'db_name': 'SNOMED CT', 'db_id': '258351006'}          -1   \n",
       "5  {'db_name': 'SNOMED CT', 'db_id': '258353009'}           0   \n",
       "6  {'db_name': 'SNOMED CT', 'db_id': '302548004'}           5   \n",
       "7   {'db_name': 'SNOMED CT', 'db_id': '49928004'}          -1   \n",
       "8  {'db_name': 'SNOMED CT', 'db_id': '113341005'}          -1   \n",
       "9  {'db_name': 'SNOMED CT', 'db_id': '301898006'}          -1   \n",
       "\n",
       "   pred_index_score   pred_top  pred_top_score document_id  \n",
       "0          1.000000  258695005        1.000000  Albers.txt  \n",
       "1          1.000000  258695005        1.000000  Albers.txt  \n",
       "2          1.000000  258695005        1.000000  Albers.txt  \n",
       "3          0.166484   48333001        0.188725  Albers.txt  \n",
       "4               NaN  258695005        1.000000  Albers.txt  \n",
       "5          0.141101  258353009        0.141101  Albers.txt  \n",
       "6          0.114893   69536005        0.223935  Albers.txt  \n",
       "7               NaN   45048000        0.241662  Albers.txt  \n",
       "8               NaN  258695005        1.000000  Albers.txt  \n",
       "9               NaN   34763001        0.138157  Albers.txt  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ea_df_rr_ws = error_analysis(ground_truth, reranked, tasks=['nen'])\n",
    "ea_df_rr_ws.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60390bfb-8c81-4a97-982c-259e7a2b88e2",
   "metadata": {},
   "source": [
    "# TODO: Fine-tuned Re-ranker?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d02329c4-125a-46e9-a2c2-32835dd20230",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train a re-ranker once annotated data (e.g., for GraSCCo) are available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edc00a25-a88a-4498-883d-1763301dd25d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:xmen]",
   "language": "python",
   "name": "conda-env-xmen-py"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}