examples/02_spaCy_German.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a3e8444-54ec-41a2-aaef-3849a22cde57",
   "metadata": {},
   "source": [
    "# Combinining spaCy NER models with xMEN for German Clinical Entity Linking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eace0372-6610-4a03-bab6-32e186911bac",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68020a8a-9c53-4f75-a08e-5078ec5b7f35",
   "metadata": {},
   "source": [
    "### Download NER Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "50943603-e8d3-4157-a675-0d9f41f5479c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Consider using `hf_transfer` for faster downloads. This solution comes with some limitations. See https://huggingface.co/docs/huggingface_hub/hf_transfer for more details.\n",
      "../local_files/de_ggponc_medbertde-1.0.0-py3-none-any.whl\n"
     ]
    }
   ],
   "source": [
    "!huggingface-cli download phlobo/de_ggponc_medbertde de_ggponc_medbertde-1.0.0-py3-none-any.whl --local-dir ../local_files\n",
    "!pip install -q ../local_files/de_ggponc_medbertde-1.0.0-py3-none-any.whl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f216b159-7112-473e-a09e-6eff33fa7a34",
   "metadata": {},
   "source": [
    "### Prepare dicts and index\n",
    "\n",
    "`xmen dict conf/ggponc.yaml`\n",
    "\n",
    "`xmen index conf/ggponc.yaml --all --overwrite`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "407e49a2-7a40-4982-bf1b-d15de505d039",
   "metadata": {},
   "source": [
    "## Run spaCy NER Model on Sample Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "521b389e-ce31-4e1f-9e37-29986d475a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load('de_ggponc_medbertde')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8049cf19-4a70-4bac-af9a-416c4e7a60e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Cetuximab ist ein monoklonaler Antikörper, der gegen den epidermalen Wachstumsfaktorrezeptor (EGFR) gerichtet ist unddient zur Therapie des fortgeschrittenen kolorektalen Karzinoms zusammen mit Irinotecan oder in Kombination mit FOLFOX bzw. allein nach Versagen einer Behandlung mit Oxaliplatin und Irinotecan.',\n",
       " 'Die HPV-Diagnostik hat beim Plattenepithelkarzinom der Mundhöhle keinen validen Nutzen als prognostischer Faktor.']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentences = [\n",
    "    \"Cetuximab ist ein monoklonaler Antikörper, der gegen den epidermalen Wachstumsfaktorrezeptor (EGFR) gerichtet ist und\" \\\n",
    "       \"dient zur Therapie des fortgeschrittenen kolorektalen Karzinoms zusammen mit Irinotecan oder in Kombination mit FOLFOX bzw. \" \\\n",
    "       \"allein nach Versagen einer Behandlung mit Oxaliplatin und Irinotecan.\",\n",
    "    \"Die HPV-Diagnostik hat beim Plattenepithelkarzinom der Mundhöhle keinen validen Nutzen als prognostischer Faktor.\"\n",
    "]\n",
    "sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5feca36d-cff3-4e02-858e-8d46c0c3d579",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = list(nlp.pipe(sentences))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a663fcf-c369-4e0c-9b90-c51fecba6019",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bd0ab2ae-bd78-4c4f-bec7-e8fc119690e1",
   "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>mention</th>\n",
       "      <th>class</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>monoklonaler Antikörper</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>epidermalen Wachstumsfaktorrezeptor</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EGFR</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Therapie des fortgeschrittenen kolorektalen Ka...</td>\n",
       "      <td>Therapeutic</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fortgeschrittenen kolorektalen Karzinoms</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>FOLFOX</td>\n",
       "      <td>Therapeutic</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Versagen einer Behandlung</td>\n",
       "      <td>Other_Finding</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Behandlung mit Oxaliplatin und Irinotecan</td>\n",
       "      <td>Therapeutic</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>HPV-Diagnostik</td>\n",
       "      <td>Diagnostic</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Plattenepithelkarzinom der Mundhöhle</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              mention  \\\n",
       "0                                           Cetuximab   \n",
       "1                             monoklonaler Antikörper   \n",
       "2                 epidermalen Wachstumsfaktorrezeptor   \n",
       "3                                                EGFR   \n",
       "4   Therapie des fortgeschrittenen kolorektalen Ka...   \n",
       "5            fortgeschrittenen kolorektalen Karzinoms   \n",
       "6                                          Irinotecan   \n",
       "7                                              FOLFOX   \n",
       "8                           Versagen einer Behandlung   \n",
       "9           Behandlung mit Oxaliplatin und Irinotecan   \n",
       "10                                        Oxaliplatin   \n",
       "11                                         Irinotecan   \n",
       "12                                     HPV-Diagnostik   \n",
       "13               Plattenepithelkarzinom der Mundhöhle   \n",
       "\n",
       "                         class  \n",
       "0                Clinical_Drug  \n",
       "1                Clinical_Drug  \n",
       "2   Nutrient_or_Body_Substance  \n",
       "3   Nutrient_or_Body_Substance  \n",
       "4                  Therapeutic  \n",
       "5       Diagnosis_or_Pathology  \n",
       "6                Clinical_Drug  \n",
       "7                  Therapeutic  \n",
       "8                Other_Finding  \n",
       "9                  Therapeutic  \n",
       "10               Clinical_Drug  \n",
       "11               Clinical_Drug  \n",
       "12                  Diagnostic  \n",
       "13      Diagnosis_or_Pathology  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ents = []\n",
    "for d in docs:\n",
    "    for span in sorted(d.spans['entities'], key=lambda s: s.start):\n",
    "        ents.append({'mention' : span.text, 'class' : span.label_})\n",
    "pd.DataFrame(ents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ae82793-8135-41a2-b6f5-75e39e1e57b1",
   "metadata": {},
   "source": [
    "## Candidate Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f5b0add5-9fb7-4eb4-9495-91ee5a5c95ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.data import from_spacy\n",
    "from xmen.linkers import SapBERTLinker, TFIDFNGramLinker, EnsembleLinker\n",
    "from xmen import load_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "46831d66-e521-4225-9b92-a40d7b7582e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = from_spacy(docs, span_key='entities')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5abb3126-ba45-47b1-a1a6-912652fefa3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['id', 'document_id', 'passages', 'entities', 'coreferences', 'relations', 'events', 'corpus_id', 'lang'],\n",
       "    num_rows: 2\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "286a07dc-e5a3-42af-89f7-b48e035f0d25",
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = load_config('../examples/conf/ggponc.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e3abe18e-4774-4832-972c-5ad80c305f8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/Florian.Borchert/miniconda3/envs/xmen_notebooks/lib/python3.10/site-packages/sklearn/base.py:348: InconsistentVersionWarning: Trying to unpickle estimator TfidfTransformer from version 1.2.2 when using version 1.3.1. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
      "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
      "  warnings.warn(\n",
      "/home/Florian.Borchert/miniconda3/envs/xmen_notebooks/lib/python3.10/site-packages/sklearn/base.py:348: InconsistentVersionWarning: Trying to unpickle estimator TfidfVectorizer from version 1.2.2 when using version 1.3.1. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
      "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "ngram_linker = TFIDFNGramLinker(**conf.linker.candidate_generation.ngram)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c932a252-1fd0-4f3c-bf7d-ac20095aa440",
   "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\">[01/07/24 15:45:41] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loading hierarchical faiss index                                <a href=\"file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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[01/07/24 15:45:41]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading hierarchical faiss index                                \u001b]8;id=923288;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/sap_bert_linker.py\u001b\\\u001b[2msap_bert_linker.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=19178;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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>         <span style=\"color: #800080; text-decoration-color: #800080\">/home/Florian.Borchert/.cache/xmen/ggponc/index/sapbert/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">embed_fais</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">s_hier.pickle</span>                                                      <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=506965;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=220384;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#64\u001b\\\u001b[2m64\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[35m/home/Florian.Borchert/.cache/xmen/ggponc/index/sapbert/\u001b[0m\u001b[95membed_fais\u001b[0m \u001b[2m                   \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[95ms_hier.pickle\u001b[0m                                                      \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\">[01/07/24 15:45:46] </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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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\">2906321</span>                                                       <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[01/07/24 15:45:46]\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=738004;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=185609;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#66\u001b\\\u001b[2m66\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         size \u001b[1;36m2906321\u001b[0m                                                       \u001b[2m                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "SapBERTLinker.clear()\n",
    "sap_linker = SapBERTLinker(cuda=False, **conf.linker.candidate_generation.sapbert)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ec6cd089-3a63-45e8-9178-86999e7f0441",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "84f76984f37d4c4a904fd8012fcca09b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "linker = EnsembleLinker()\n",
    "linker.add_linker('ngram', ngram_linker, k=conf.linker.candidate_generation.ngram.k, threshold=0.9)\n",
    "linker.add_linker('sap', sap_linker, k=conf.linker.candidate_generation.sapbert.k, threshold=0.8)\n",
    "\n",
    "prediction = linker.predict_batch(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3689a61-2ee1-4c81-afc7-6f6751654020",
   "metadata": {},
   "source": [
    "### Semantic Type Filtering\n",
    "\n",
    "We filter the generated output to make sure the semantic type of the predicted concepts actually matches the semantic class of the named entity.\n",
    "\n",
    "As the GGPONC entity classes are based on SNOMED CT top level concepts, while we link against UMLS CUIS, we provide a mapping of GGPONC enitity types to UMLS TUIs in `ggponc2tui.tsv`\n",
    "\n",
    "Semantic Type Filtering is particularly useful for ambiguous abbreviations (e.g., \"EGFR\" as in the example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "515adb34-e2de-4f7f-b133-41ddf8ca421d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.kb import load_kb\n",
    "from xmen.data import SemanticTypeFilter\n",
    "import pandas as pd\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "29f28a60-d5d4-4ee3-bedb-55e71d15fe84",
   "metadata": {},
   "outputs": [],
   "source": [
    "kb = load_kb(Path(conf.cache_dir) / 'ggponc' / 'ggponc.jsonl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "df540489-bdb6-4e60-a4bf-59860861d284",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.data import SemanticTypeFilter\n",
    "\n",
    "type2tui = pd.read_csv('ggponc_tuis.csv').groupby('class')['tui'].apply(list).to_dict()\n",
    "type_filter = SemanticTypeFilter(type2tui, kb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "14ebda62-a868-4bc9-8c84-a0580935b0c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ba939fdda2e40a78cf02de6f7b5d580",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "filtered_prediction = type_filter.transform_batch(prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "308b40e7-d97a-4b58-b933-ecc8083fcc97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['EGFR']\n",
      "{'db_id': 'C3811844', 'db_name': 'UMLS', 'score': 1.0, 'predicted_by': ['ngram', 'sap']}\n",
      "CUI: C3811844, Name: Geschaetzte glomerulaere Filtrationsrate\n",
      "Definition: A laboratory test that estimates kidney function. It is calculated using an individual's serum creatinine measurement, age, gender, and race. Actual results are reported when the estimated glomerular filtration rate is less than 60 ml/min.\n",
      "TUI(s): T059\n",
      "Aliases: (total: 3): \n",
      "\t Estimated Glomerular Filtration Rate, eGFR, Estimated glomerular filtration rate\n"
     ]
    }
   ],
   "source": [
    "# Before Filtering\n",
    "entity = prediction[0]['entities'][3]\n",
    "print(entity['text'])\n",
    "print(entity['normalized'][0])\n",
    "print(kb.cui_to_entity['C3811844'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bc358d6c-c1a6-4b8a-9839-3a527c3f2a67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['EGFR']\n",
      "{'db_id': 'C1739039', 'db_name': 'UMLS', 'score': 1.0, 'predicted_by': ['ngram', 'sap']}\n",
      "CUI: C1739039, Name: EC 2.7.10.1\n",
      "Definition: The protein found on the surface of some cells and to which epidermal growth factor binds, causing the cells to divide. It is found at abnormally high levels on the surface of many types of cancer cells, so these cells may divide excessively in the presence of epidermal growth factor.\n",
      "TUI(s): T116, T192\n",
      "Aliases (abbreviated, total: 20): \n",
      "\t ERBB protein, human, EGF Receptor, EGFR protein, human, EGFR, Erb-B2 Receptor Tyrosine Kinase 1, HER-1, Epidermal Growth Factor Receptor, ERRP protein, human, Urogastrone Receptor, HER1 protein, human\n"
     ]
    }
   ],
   "source": [
    "# After Filtering\n",
    "entity = filtered_prediction[0]['entities'][3]\n",
    "print(entity['text'])\n",
    "print(entity['normalized'][0])\n",
    "print(kb.cui_to_entity['C1739039'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c02ef4fd-dc7f-487d-b059-8643b3886ba5",
   "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>mention</th>\n",
       "      <th>class</th>\n",
       "      <th>cui</th>\n",
       "      <th>canonical name</th>\n",
       "      <th>linked by</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0995188</td>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>monoklonaler Antikörper</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0003250</td>\n",
       "      <td>Antikörper, monoklonale</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>0.982318</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>epidermalen Wachstumsfaktorrezeptor</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>C3812393</td>\n",
       "      <td>ErbB-Rezeptoren</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>0.957236</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EGFR</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>C1739039</td>\n",
       "      <td>EC 2.7.10.1</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Therapie des fortgeschrittenen kolorektalen Ka...</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>Not linkable</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fortgeschrittenen kolorektalen Karzinoms</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>C4721579</td>\n",
       "      <td>Kolorektalkarzinom mit Metastasen</td>\n",
       "      <td>[sap]</td>\n",
       "      <td>0.843411</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0123931</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>FOLFOX</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0309154</td>\n",
       "      <td>FUROX (product)</td>\n",
       "      <td>[sap]</td>\n",
       "      <td>0.842153</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Versagen einer Behandlung</td>\n",
       "      <td>Other_Finding</td>\n",
       "      <td>C0162643</td>\n",
       "      <td>Behandlungsfehler</td>\n",
       "      <td>[sap]</td>\n",
       "      <td>0.936106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Behandlung mit Oxaliplatin und Irinotecan</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0796324</td>\n",
       "      <td>IROX Regimen</td>\n",
       "      <td>[sap]</td>\n",
       "      <td>0.879975</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0069717</td>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0123931</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sap]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>HPV-Diagnostik</td>\n",
       "      <td>Diagnostic</td>\n",
       "      <td>C3839551</td>\n",
       "      <td>Human papilloma virus screening</td>\n",
       "      <td>[sap]</td>\n",
       "      <td>0.913146</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Plattenepithelkarzinom der Mundhöhle</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>C0585362</td>\n",
       "      <td>Plattenepithelkarzinom der Mundhoehle</td>\n",
       "      <td>[sap]</td>\n",
       "      <td>0.987126</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              mention  \\\n",
       "0                                           Cetuximab   \n",
       "1                             monoklonaler Antikörper   \n",
       "2                 epidermalen Wachstumsfaktorrezeptor   \n",
       "3                                                EGFR   \n",
       "4   Therapie des fortgeschrittenen kolorektalen Ka...   \n",
       "5            fortgeschrittenen kolorektalen Karzinoms   \n",
       "6                                          Irinotecan   \n",
       "7                                              FOLFOX   \n",
       "8                           Versagen einer Behandlung   \n",
       "9           Behandlung mit Oxaliplatin und Irinotecan   \n",
       "10                                        Oxaliplatin   \n",
       "11                                         Irinotecan   \n",
       "12                                     HPV-Diagnostik   \n",
       "13               Plattenepithelkarzinom der Mundhöhle   \n",
       "\n",
       "                         class           cui  \\\n",
       "0                Clinical_Drug      C0995188   \n",
       "1                Clinical_Drug      C0003250   \n",
       "2   Nutrient_or_Body_Substance      C3812393   \n",
       "3   Nutrient_or_Body_Substance      C1739039   \n",
       "4                  Therapeutic  Not linkable   \n",
       "5       Diagnosis_or_Pathology      C4721579   \n",
       "6                Clinical_Drug      C0123931   \n",
       "7                  Therapeutic      C0309154   \n",
       "8                Other_Finding      C0162643   \n",
       "9                  Therapeutic      C0796324   \n",
       "10               Clinical_Drug      C0069717   \n",
       "11               Clinical_Drug      C0123931   \n",
       "12                  Diagnostic      C3839551   \n",
       "13      Diagnosis_or_Pathology      C0585362   \n",
       "\n",
       "                           canonical name     linked by     score  \n",
       "0                               Cetuximab  [ngram, sap]  1.000000  \n",
       "1                 Antikörper, monoklonale  [ngram, sap]  0.982318  \n",
       "2                         ErbB-Rezeptoren  [ngram, sap]  0.957236  \n",
       "3                             EC 2.7.10.1  [ngram, sap]  1.000000  \n",
       "4                                     NaN           NaN       NaN  \n",
       "5       Kolorektalkarzinom mit Metastasen         [sap]  0.843411  \n",
       "6                              Irinotecan  [ngram, sap]  1.000000  \n",
       "7                         FUROX (product)         [sap]  0.842153  \n",
       "8                       Behandlungsfehler         [sap]  0.936106  \n",
       "9                            IROX Regimen         [sap]  0.879975  \n",
       "10                            Oxaliplatin  [ngram, sap]  1.000000  \n",
       "11                             Irinotecan  [ngram, sap]  1.000000  \n",
       "12        Human papilloma virus screening         [sap]  0.913146  \n",
       "13  Plattenepithelkarzinom der Mundhoehle         [sap]  0.987126  "
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from util import get_dataframe\n",
    "get_dataframe(filtered_prediction, kb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b8847b5-ff8f-40de-8a18-f2c2485daf7e",
   "metadata": {},
   "source": [
    "## Re-Ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9b7b5c5f-e617-4975-822a-fa6bce38397a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/Florian.Borchert/miniconda3/envs/xmen_notebooks/lib/python3.10/site-packages/sklearn/base.py:348: InconsistentVersionWarning: Trying to unpickle estimator TfidfTransformer from version 1.2.2 when using version 1.3.1. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
      "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
      "  warnings.warn(\n",
      "/home/Florian.Borchert/miniconda3/envs/xmen_notebooks/lib/python3.10/site-packages/sklearn/base.py:348: InconsistentVersionWarning: Trying to unpickle estimator TfidfVectorizer from version 1.2.2 when using version 1.3.1. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
      "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "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\">[01/07/24 15:46:19] </span><span style=\"color: #000080; text-decoration-color: #000080\">INFO    </span> Loading hierarchical faiss index                                <a href=\"file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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[01/07/24 15:46:19]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading hierarchical faiss index                                \u001b]8;id=428661;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/sap_bert_linker.py\u001b\\\u001b[2msap_bert_linker.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=669910;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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>         <span style=\"color: #800080; text-decoration-color: #800080\">/home/Florian.Borchert/.cache/xmen/ggponc/index/sapbert/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">embed_fais</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">s_hier.pickle</span>                                                      <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=26558;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=576197;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#64\u001b\\\u001b[2m64\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[35m/home/Florian.Borchert/.cache/xmen/ggponc/index/sapbert/\u001b[0m\u001b[95membed_fais\u001b[0m \u001b[2m                   \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[95ms_hier.pickle\u001b[0m                                                      \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\">[01/07/24 15:46:23] </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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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:///mnt/nfs/home/Florian.Borchert/workspace/xmen/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\">2906321</span>                                                       <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[01/07/24 15:46:23]\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=783911;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=540849;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#66\u001b\\\u001b[2m66\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         size \u001b[1;36m2906321\u001b[0m                                                       \u001b[2m                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from xmen.linkers import default_ensemble\n",
    "linker_no_thresh = default_ensemble(Path(conf.linker.candidate_generation.ngram.index_base_path).parent, cuda=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "411c0814-5fc5-4b70-ab96-ef07833cbf08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aff39c33fd664530b77c54d10aafd1ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3cf871c71924990af44012caf67d1e5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "candidates = type_filter.transform_batch(linker_no_thresh.predict_batch(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8185c65f-5645-4c38-a11a-a32a4f24a2a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.reranking import CrossEncoderReranker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "07d41590-f991-474a-9109-9b75d3184266",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context length: 128\n",
      "Use NIL values: True\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c63a948ec5094b889319718786deed34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99a147c50db341268f76bc3744a58955",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db70755198584202872f64b700f400a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67fef449ec1e416ab904c7c4d04161fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ce_candidates = CrossEncoderReranker.prepare_data(candidates, None, kb, k=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6613f881-9323-42c7-a63c-710c6a1d758c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr = CrossEncoderReranker.load(\"phlobo/xmen-de-ce-medmentions\", device=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3af8dfbf-14b1-4a47-b76b-14868cde55e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4b418017fc724e7c9e9aa8925e152564",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc9435aeb6c6491da2ab23f569a6aa2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf493af6ff4a4ef5b211ee3e677e0d9e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3877c04868c34abe84b3356736ce2270",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "reranked = rr.rerank_batch(candidates, ce_candidates, k=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8bdfb06e-a38b-431d-b82a-c515e6a00078",
   "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>mention</th>\n",
       "      <th>class</th>\n",
       "      <th>cui</th>\n",
       "      <th>canonical name</th>\n",
       "      <th>linked by</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0995188</td>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>monoklonaler Antikörper</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0003250</td>\n",
       "      <td>Antikörper, monoklonale</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.982318</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>epidermalen Wachstumsfaktorrezeptor</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>C3812393</td>\n",
       "      <td>ErbB-Rezeptoren</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.957236</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EGFR</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>C1739039</td>\n",
       "      <td>EC 2.7.10.1</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Therapie des fortgeschrittenen kolorektalen Ka...</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0281190</td>\n",
       "      <td>Prevention of Colorectal Cancer</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.736036</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fortgeschrittenen kolorektalen Karzinoms</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>C4721579</td>\n",
       "      <td>Kolorektalkarzinom mit Metastasen</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.843411</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0123931</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>FOLFOX</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0309154</td>\n",
       "      <td>FUROX (product)</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.842153</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Versagen einer Behandlung</td>\n",
       "      <td>Other_Finding</td>\n",
       "      <td>C0162643</td>\n",
       "      <td>Behandlungsfehler</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.936106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Behandlung mit Oxaliplatin und Irinotecan</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0796324</td>\n",
       "      <td>IROX Regimen</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.879975</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0069717</td>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0123931</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>HPV-Diagnostik</td>\n",
       "      <td>Diagnostic</td>\n",
       "      <td>C3839551</td>\n",
       "      <td>Human papilloma virus screening</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.913146</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Plattenepithelkarzinom der Mundhöhle</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>C0585362</td>\n",
       "      <td>Plattenepithelkarzinom der Mundhoehle</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.987126</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              mention  \\\n",
       "0                                           Cetuximab   \n",
       "1                             monoklonaler Antikörper   \n",
       "2                 epidermalen Wachstumsfaktorrezeptor   \n",
       "3                                                EGFR   \n",
       "4   Therapie des fortgeschrittenen kolorektalen Ka...   \n",
       "5            fortgeschrittenen kolorektalen Karzinoms   \n",
       "6                                          Irinotecan   \n",
       "7                                              FOLFOX   \n",
       "8                           Versagen einer Behandlung   \n",
       "9           Behandlung mit Oxaliplatin und Irinotecan   \n",
       "10                                        Oxaliplatin   \n",
       "11                                         Irinotecan   \n",
       "12                                     HPV-Diagnostik   \n",
       "13               Plattenepithelkarzinom der Mundhöhle   \n",
       "\n",
       "                         class       cui  \\\n",
       "0                Clinical_Drug  C0995188   \n",
       "1                Clinical_Drug  C0003250   \n",
       "2   Nutrient_or_Body_Substance  C3812393   \n",
       "3   Nutrient_or_Body_Substance  C1739039   \n",
       "4                  Therapeutic  C0281190   \n",
       "5       Diagnosis_or_Pathology  C4721579   \n",
       "6                Clinical_Drug  C0123931   \n",
       "7                  Therapeutic  C0309154   \n",
       "8                Other_Finding  C0162643   \n",
       "9                  Therapeutic  C0796324   \n",
       "10               Clinical_Drug  C0069717   \n",
       "11               Clinical_Drug  C0123931   \n",
       "12                  Diagnostic  C3839551   \n",
       "13      Diagnosis_or_Pathology  C0585362   \n",
       "\n",
       "                           canonical name         linked by     score  \n",
       "0                               Cetuximab  [ngram, sapbert]  1.000000  \n",
       "1                 Antikörper, monoklonale  [ngram, sapbert]  0.982318  \n",
       "2                         ErbB-Rezeptoren  [ngram, sapbert]  0.957236  \n",
       "3                             EC 2.7.10.1  [ngram, sapbert]  1.000000  \n",
       "4         Prevention of Colorectal Cancer         [sapbert]  0.736036  \n",
       "5       Kolorektalkarzinom mit Metastasen  [ngram, sapbert]  0.843411  \n",
       "6                              Irinotecan  [ngram, sapbert]  1.000000  \n",
       "7                         FUROX (product)         [sapbert]  0.842153  \n",
       "8                       Behandlungsfehler  [ngram, sapbert]  0.936106  \n",
       "9                            IROX Regimen         [sapbert]  0.879975  \n",
       "10                            Oxaliplatin  [ngram, sapbert]  1.000000  \n",
       "11                             Irinotecan  [ngram, sapbert]  1.000000  \n",
       "12        Human papilloma virus screening         [sapbert]  0.913146  \n",
       "13  Plattenepithelkarzinom der Mundhoehle  [ngram, sapbert]  0.987126  "
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Before Re-ranking\n",
    "get_dataframe(candidates, kb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0f95ab2d-5f3b-432c-9214-64316b3f1e99",
   "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>mention</th>\n",
       "      <th>class</th>\n",
       "      <th>cui</th>\n",
       "      <th>canonical name</th>\n",
       "      <th>linked by</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0995188</td>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.154751</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>monoklonaler Antikörper</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0003250</td>\n",
       "      <td>Antikörper, monoklonale</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.159235</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>epidermalen Wachstumsfaktorrezeptor</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>C3812393</td>\n",
       "      <td>ErbB-Rezeptoren</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.081309</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EGFR</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>C1739039</td>\n",
       "      <td>EC 2.7.10.1</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.079179</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Therapie des fortgeschrittenen kolorektalen Ka...</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C4763871</td>\n",
       "      <td>Colorectal Cancer Surgery</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.069189</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fortgeschrittenen kolorektalen Karzinoms</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>C0009402</td>\n",
       "      <td>Kolorektales Karzinom</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.076059</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0123931</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.167424</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>FOLFOX</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0392943</td>\n",
       "      <td>Fluorouracil/Leucovorin Calcium/Oxaliplatin</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.098444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Versagen einer Behandlung</td>\n",
       "      <td>Other_Finding</td>\n",
       "      <td>C0162643</td>\n",
       "      <td>Behandlungsfehler</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.145101</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Behandlung mit Oxaliplatin und Irinotecan</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>C0796324</td>\n",
       "      <td>IROX Regimen</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.077630</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0069717</td>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.156690</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>C0123931</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.173561</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>HPV-Diagnostik</td>\n",
       "      <td>Diagnostic</td>\n",
       "      <td>C1610047</td>\n",
       "      <td>Human papilloma virus test</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.073448</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Plattenepithelkarzinom der Mundhöhle</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>C0585362</td>\n",
       "      <td>Plattenepithelkarzinom der Mundhoehle</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.148823</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              mention  \\\n",
       "0                                           Cetuximab   \n",
       "1                             monoklonaler Antikörper   \n",
       "2                 epidermalen Wachstumsfaktorrezeptor   \n",
       "3                                                EGFR   \n",
       "4   Therapie des fortgeschrittenen kolorektalen Ka...   \n",
       "5            fortgeschrittenen kolorektalen Karzinoms   \n",
       "6                                          Irinotecan   \n",
       "7                                              FOLFOX   \n",
       "8                           Versagen einer Behandlung   \n",
       "9           Behandlung mit Oxaliplatin und Irinotecan   \n",
       "10                                        Oxaliplatin   \n",
       "11                                         Irinotecan   \n",
       "12                                     HPV-Diagnostik   \n",
       "13               Plattenepithelkarzinom der Mundhöhle   \n",
       "\n",
       "                         class       cui  \\\n",
       "0                Clinical_Drug  C0995188   \n",
       "1                Clinical_Drug  C0003250   \n",
       "2   Nutrient_or_Body_Substance  C3812393   \n",
       "3   Nutrient_or_Body_Substance  C1739039   \n",
       "4                  Therapeutic  C4763871   \n",
       "5       Diagnosis_or_Pathology  C0009402   \n",
       "6                Clinical_Drug  C0123931   \n",
       "7                  Therapeutic  C0392943   \n",
       "8                Other_Finding  C0162643   \n",
       "9                  Therapeutic  C0796324   \n",
       "10               Clinical_Drug  C0069717   \n",
       "11               Clinical_Drug  C0123931   \n",
       "12                  Diagnostic  C1610047   \n",
       "13      Diagnosis_or_Pathology  C0585362   \n",
       "\n",
       "                                 canonical name         linked by     score  \n",
       "0                                     Cetuximab  [ngram, sapbert]  0.154751  \n",
       "1                       Antikörper, monoklonale  [ngram, sapbert]  0.159235  \n",
       "2                               ErbB-Rezeptoren  [ngram, sapbert]  0.081309  \n",
       "3                                   EC 2.7.10.1  [ngram, sapbert]  0.079179  \n",
       "4                     Colorectal Cancer Surgery         [sapbert]  0.069189  \n",
       "5                         Kolorektales Karzinom  [ngram, sapbert]  0.076059  \n",
       "6                                    Irinotecan  [ngram, sapbert]  0.167424  \n",
       "7   Fluorouracil/Leucovorin Calcium/Oxaliplatin  [ngram, sapbert]  0.098444  \n",
       "8                             Behandlungsfehler  [ngram, sapbert]  0.145101  \n",
       "9                                  IROX Regimen         [sapbert]  0.077630  \n",
       "10                                  Oxaliplatin  [ngram, sapbert]  0.156690  \n",
       "11                                   Irinotecan  [ngram, sapbert]  0.173561  \n",
       "12                   Human papilloma virus test         [sapbert]  0.073448  \n",
       "13        Plattenepithelkarzinom der Mundhoehle  [ngram, sapbert]  0.148823  "
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# After Re-ranking\n",
    "get_dataframe(reranked, kb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeda34f6-da40-45a7-99e6-36fbe9ae1c43",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:xmen_notebooks]",
   "language": "python",
   "name": "conda-env-xmen_notebooks-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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}