examples/03_SNOMED_Linking_German.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a3e8444-54ec-41a2-aaef-3849a22cde57",
   "metadata": {},
   "source": [
    "# Linking Entities in German Medical Text to SNOMED CT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eace0372-6610-4a03-bab6-32e186911bac",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "### 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/snomed_german.yaml --code dicts/umls_source.py`\n",
    "\n",
    "`xmen index conf/snomed_german.yaml --all --overwrite`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "407e49a2-7a40-4982-bf1b-d15de505d039",
   "metadata": {},
   "source": [
    "## Entity Tagging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "521b389e-ce31-4e1f-9e37-29986d475a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load('de_ggponc_medbertde')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 3,
     "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": 4,
   "id": "5feca36d-cff3-4e02-858e-8d46c0c3d579",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = list(nlp.pipe(sentences))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a663fcf-c369-4e0c-9b90-c51fecba6019",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 6,
     "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": 7,
   "id": "f5b0add5-9fb7-4eb4-9495-91ee5a5c95ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "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": 8,
   "id": "46831d66-e521-4225-9b92-a40d7b7582e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = from_spacy(docs, span_key='entities')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "286a07dc-e5a3-42af-89f7-b48e035f0d25",
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = load_config('../examples/conf/snomed_german.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ec6cd089-3a63-45e8-9178-86999e7f0441",
   "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/09/24 16:31:37] </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/09/24 16:31:37]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading hierarchical faiss index                                \u001b]8;id=728264;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=957841;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/snomed_german/index/sapbert/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">emb</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\">ed_faiss_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=170528;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=461671;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/snomed_german/index/sapbert/\u001b[0m\u001b[95memb\u001b[0m \u001b[2m                   \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[95med_faiss_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/09/24 16:31:40] </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\">1967771</span>                                                       <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[01/09/24 16:31:40]\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=355658;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=291906;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;36m1967771\u001b[0m                                                       \u001b[2m                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7068e4ad1fc42a1959b309292787250",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from xmen.linkers import default_ensemble\n",
    "linker = default_ensemble(Path(conf.linker.candidate_generation.ngram.index_base_path).parent, cuda=False)\n",
    "\n",
    "candidates = linker.predict_batch(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3728665d-03ae-47a6-82d6-3cf906b72c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen import load_kb\n",
    "\n",
    "kb = load_kb(Path(conf.cache_dir) / 'snomed_german' / 'snomed_german.jsonl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c3a659c4-0371-4b6f-9544-8b1fc6615c61",
   "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>409401002</td>\n",
       "      <td>Cetuximab-containing product</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>49616005</td>\n",
       "      <td>Monoclonal antibody</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>86960007</td>\n",
       "      <td>Epidermal growth factor-urogastrone receptor (...</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.937485</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EGFR</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>86960007</td>\n",
       "      <td>Epidermal growth factor-urogastrone receptor (...</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>1217692004</td>\n",
       "      <td>Metastasis from malignant neoplasm of colon an...</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.696480</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fortgeschrittenen kolorektalen Karzinoms</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>1217692004</td>\n",
       "      <td>Metastasis from malignant neoplasm of colon an...</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>372538008</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>699297004</td>\n",
       "      <td>Ohdo syndrome, Maat-Kievit-Brunner type</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.812118</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Versagen einer Behandlung</td>\n",
       "      <td>Other_Finding</td>\n",
       "      <td>7058009</td>\n",
       "      <td>Noncompliance with treatment (finding)</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.881669</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Behandlung mit Oxaliplatin und Irinotecan</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>447053005</td>\n",
       "      <td>Oxaliplatin desensitization therapy</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.673975</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>395814003</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>372538008</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>700152009</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>307502000</td>\n",
       "      <td>Squamous cell carcinoma of mouth</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   409401002   \n",
       "1                Clinical_Drug    49616005   \n",
       "2   Nutrient_or_Body_Substance    86960007   \n",
       "3   Nutrient_or_Body_Substance    86960007   \n",
       "4                  Therapeutic  1217692004   \n",
       "5       Diagnosis_or_Pathology  1217692004   \n",
       "6                Clinical_Drug   372538008   \n",
       "7                  Therapeutic   699297004   \n",
       "8                Other_Finding     7058009   \n",
       "9                  Therapeutic   447053005   \n",
       "10               Clinical_Drug   395814003   \n",
       "11               Clinical_Drug   372538008   \n",
       "12                  Diagnostic   700152009   \n",
       "13      Diagnosis_or_Pathology   307502000   \n",
       "\n",
       "                                       canonical name         linked by  \\\n",
       "0                        Cetuximab-containing product  [ngram, sapbert]   \n",
       "1                                 Monoclonal antibody  [ngram, sapbert]   \n",
       "2   Epidermal growth factor-urogastrone receptor (...         [sapbert]   \n",
       "3   Epidermal growth factor-urogastrone receptor (...  [ngram, sapbert]   \n",
       "4   Metastasis from malignant neoplasm of colon an...  [ngram, sapbert]   \n",
       "5   Metastasis from malignant neoplasm of colon an...  [ngram, sapbert]   \n",
       "6                                          Irinotecan  [ngram, sapbert]   \n",
       "7             Ohdo syndrome, Maat-Kievit-Brunner type         [sapbert]   \n",
       "8              Noncompliance with treatment (finding)  [ngram, sapbert]   \n",
       "9                 Oxaliplatin desensitization therapy  [ngram, sapbert]   \n",
       "10                                        Oxaliplatin  [ngram, sapbert]   \n",
       "11                                         Irinotecan  [ngram, sapbert]   \n",
       "12                    Human papilloma virus screening         [sapbert]   \n",
       "13                   Squamous cell carcinoma of mouth  [ngram, sapbert]   \n",
       "\n",
       "       score  \n",
       "0   1.000000  \n",
       "1   0.982318  \n",
       "2   0.937485  \n",
       "3   1.000000  \n",
       "4   0.696480  \n",
       "5   0.843411  \n",
       "6   1.000000  \n",
       "7   0.812118  \n",
       "8   0.881669  \n",
       "9   0.673975  \n",
       "10  1.000000  \n",
       "11  1.000000  \n",
       "12  0.913146  \n",
       "13  0.987126  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from util import get_dataframe\n",
    "get_dataframe(candidates, kb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b8847b5-ff8f-40de-8a18-f2c2485daf7e",
   "metadata": {},
   "source": [
    "## Re-Ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8185c65f-5645-4c38-a11a-a32a4f24a2a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.reranking import CrossEncoderReranker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "07d41590-f991-474a-9109-9b75d3184266",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context length: 128\n",
      "Use NIL values: False\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a01966a82ade44ba92a3c8b974b2b889",
       "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": "5611f72062604097bebd030fbbef9478",
       "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": "fb0b6147f70d4938b9de16fafadc7915",
       "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": "0beee863844b4242b8354226cf0c4841",
       "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, use_nil=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6613f881-9323-42c7-a63c-710c6a1d758c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr = CrossEncoderReranker.load(\"phlobo/xmen-de-ce-medmentions\", device=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3af8dfbf-14b1-4a47-b76b-14868cde55e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae4eded3c1d74f2db4b685110dbbb805",
       "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": "63c20dcf6eb048e8adf9b40f243725bf",
       "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": "1e2f4905210e45ba87efd813b77ed634",
       "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": "7b237ed60e8a428e8d257e805866ffd1",
       "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": 18,
   "id": "62128aa7-3434-4768-a36a-69cce7f2d294",
   "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>409400001</td>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.136370</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>monoklonaler Antikörper</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>49616005</td>\n",
       "      <td>Monoclonal antibody</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.145249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>epidermalen Wachstumsfaktorrezeptor</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>86960007</td>\n",
       "      <td>Epidermal growth factor-urogastrone receptor (...</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.080162</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EGFR</td>\n",
       "      <td>Nutrient_or_Body_Substance</td>\n",
       "      <td>86960007</td>\n",
       "      <td>Epidermal growth factor-urogastrone receptor (...</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.114013</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Therapie des fortgeschrittenen kolorektalen Ka...</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>1217692004</td>\n",
       "      <td>Metastasis from malignant neoplasm of colon an...</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.068034</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fortgeschrittenen kolorektalen Karzinoms</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>126837005</td>\n",
       "      <td>Neoplasm of large intestine</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.067600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>372538008</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.139657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>FOLFOX</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>461391000124102</td>\n",
       "      <td>Folinic acid-fluorouracil-oxaliplatin regimen ...</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.072592</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Versagen einer Behandlung</td>\n",
       "      <td>Other_Finding</td>\n",
       "      <td>266721009</td>\n",
       "      <td>Absent response to treatment</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.068675</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Behandlung mit Oxaliplatin und Irinotecan</td>\n",
       "      <td>Therapeutic</td>\n",
       "      <td>447053005</td>\n",
       "      <td>Oxaliplatin desensitization therapy</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.066917</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>395814003</td>\n",
       "      <td>Oxaliplatin</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.138286</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>Clinical_Drug</td>\n",
       "      <td>372538008</td>\n",
       "      <td>Irinotecan</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.140389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>HPV-Diagnostik</td>\n",
       "      <td>Diagnostic</td>\n",
       "      <td>700152009</td>\n",
       "      <td>Human papilloma virus screening</td>\n",
       "      <td>[sapbert]</td>\n",
       "      <td>0.076523</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Plattenepithelkarzinom der Mundhöhle</td>\n",
       "      <td>Diagnosis_or_Pathology</td>\n",
       "      <td>307502000</td>\n",
       "      <td>Squamous cell carcinoma of mouth</td>\n",
       "      <td>[ngram, sapbert]</td>\n",
       "      <td>0.139139</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        409400001   \n",
       "1                Clinical_Drug         49616005   \n",
       "2   Nutrient_or_Body_Substance         86960007   \n",
       "3   Nutrient_or_Body_Substance         86960007   \n",
       "4                  Therapeutic       1217692004   \n",
       "5       Diagnosis_or_Pathology        126837005   \n",
       "6                Clinical_Drug        372538008   \n",
       "7                  Therapeutic  461391000124102   \n",
       "8                Other_Finding        266721009   \n",
       "9                  Therapeutic        447053005   \n",
       "10               Clinical_Drug        395814003   \n",
       "11               Clinical_Drug        372538008   \n",
       "12                  Diagnostic        700152009   \n",
       "13      Diagnosis_or_Pathology        307502000   \n",
       "\n",
       "                                       canonical name         linked by  \\\n",
       "0                                           Cetuximab  [ngram, sapbert]   \n",
       "1                                 Monoclonal antibody  [ngram, sapbert]   \n",
       "2   Epidermal growth factor-urogastrone receptor (...         [sapbert]   \n",
       "3   Epidermal growth factor-urogastrone receptor (...  [ngram, sapbert]   \n",
       "4   Metastasis from malignant neoplasm of colon an...  [ngram, sapbert]   \n",
       "5                         Neoplasm of large intestine  [ngram, sapbert]   \n",
       "6                                          Irinotecan  [ngram, sapbert]   \n",
       "7   Folinic acid-fluorouracil-oxaliplatin regimen ...  [ngram, sapbert]   \n",
       "8                        Absent response to treatment         [sapbert]   \n",
       "9                 Oxaliplatin desensitization therapy  [ngram, sapbert]   \n",
       "10                                        Oxaliplatin  [ngram, sapbert]   \n",
       "11                                         Irinotecan  [ngram, sapbert]   \n",
       "12                    Human papilloma virus screening         [sapbert]   \n",
       "13                   Squamous cell carcinoma of mouth  [ngram, sapbert]   \n",
       "\n",
       "       score  \n",
       "0   0.136370  \n",
       "1   0.145249  \n",
       "2   0.080162  \n",
       "3   0.114013  \n",
       "4   0.068034  \n",
       "5   0.067600  \n",
       "6   0.139657  \n",
       "7   0.072592  \n",
       "8   0.068675  \n",
       "9   0.066917  \n",
       "10  0.138286  \n",
       "11  0.140389  \n",
       "12  0.076523  \n",
       "13  0.139139  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# After Re-ranking\n",
    "get_dataframe(reranked, kb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "070c2a21-fe32-4b59-a8ac-1b97d9548e76",
   "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
}