examples/01_BRONCO_German.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eb84db40-b21a-4fba-8174-a327a831ef76",
   "metadata": {},
   "source": [
    "# Preparation\n",
    "\n",
    "Make sure to download the corresponding medical ontologies to build the term dictionaries. For each, look for the 2017 version.\n",
    " - Treatments: [OPS](https://www.bfarm.de/DE/Kodiersysteme/Services/Downloads/_node.html).\n",
    " - Medications: [ATC](https://www.wido.de/publikationen-produkte/arzneimittel-klassifikation/).\n",
    " - Diagnosis: [ICD10GM](https://www.bfarm.de/DE/Kodiersysteme/Services/Downloads/_node.html). \n",
    " \n",
    "In the config file for BRONCO `conf/bronco.yaml`, modify the paths so they point the extracted folders. We assume they are located in `xmen/local_files`. Otherwise, change the path here and correct accordingly the terminal commands below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "065fa6f0-911d-4d9e-924a-faef1646e5d0",
   "metadata": {},
   "source": [
    "We can already use xMEN to prepare the term dictionaries. <u>This step only has to be performed the first time. </u> \n",
    "\n",
    "In your terminal, navigate to the xMEN root folder and run:\n",
    " - `xmen dict examples/conf/bronco.yaml --code examples/dicts/atc2017_de.py --key atc`\n",
    " - `xmen dict examples/conf/bronco.yaml --code examples/dicts/ops2017.py --key ops`\n",
    " - `xmen dict examples/conf/bronco.yaml --code examples/dicts/icd10gm2017.py --key icd10gm`\n",
    " \n",
    "Now use such dictionaries to build the indexes:\n",
    " - `xmen index examples/conf/bronco.yaml --dict ~/.cache/xmen/atc/atc.jsonl --output ~/.cache/xmen/atc/ --all`\n",
    " - `xmen index examples/conf/bronco.yaml --dict ~/.cache/xmen/ops/ops.jsonl ~/.cache/xmen/ops/ --all`\n",
    " - `xmen index examples/conf/bronco.yaml --dict ~/.cache/xmen/icd10gm/icd10gm.jsonl ~/.cache/xmen/icd10gm/ --all`\n",
    " \n",
    "Now we can load the BRONCO150 dataset using BigBIO:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "12fbc6ab-f65c-401d-bbe0-99bbcdb2cbfe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset bronco (/home/Florian.Borchert/.cache/huggingface/datasets/bigbio___bronco/bronco_bigbio_kb-fd41eed48d3255b6/1.0.0/cab8fc4a62807688cb5b36df7a24eb7f364314862c4196f6ff2db3813f2fe68b)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6cbc3e273c246288f8f431a80cdd347",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'document_id', 'passages', 'entities', 'events', 'coreferences', 'relations'],\n",
       "        num_rows: 5\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import datasets\n",
    "\n",
    "path_to_data = r\"../local_files/BRONCO150\" # paste here the path to the local data\n",
    "\n",
    "bronco = datasets.load_dataset(path = \"bigbio/bronco\", \n",
    "                               name = \"bronco_bigbio_kb\", \n",
    "                               data_dir=path_to_data)\n",
    "\n",
    "bronco"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79dc91db-df81-4a86-a601-19834659fdc7",
   "metadata": {},
   "source": [
    "Finally, we have to choose the semantic class we want to work on and reestructure the dataset in 5 folds for cross-validation, as originally intended."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "55afb5e6-4324-4ce7-b0ad-9980a6240b87",
   "metadata": {},
   "outputs": [],
   "source": [
    "label = \"MEDICATION\" # Choose here TREATMENT, MEDICATION or DIAGNOSIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3ee0ed32-c4e0-46e5-a89f-03e29b81701a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/Florian.Borchert/.cache/huggingface/datasets/bigbio___bronco/bronco_bigbio_kb-fd41eed48d3255b6/1.0.0/cab8fc4a62807688cb5b36df7a24eb7f364314862c4196f6ff2db3813f2fe68b/cache-44459d7791d9d192.arrow\n"
     ]
    }
   ],
   "source": [
    "label2dict = {\n",
    "    \"TREATMENT\": \"ops\",\n",
    "    \"MEDICATION\": \"atc\",\n",
    "    \"DIAGNOSIS\": \"icd10gm\"\n",
    "}\n",
    "\n",
    "def filter_entities(bigbio_entities, valid_entities):\n",
    "    filtered_entities = []\n",
    "    for ent in bigbio_entities:\n",
    "        if ent['type'] in valid_entities:\n",
    "            filtered_entities.append(ent)\n",
    "    return filtered_entities\n",
    "\n",
    "ds = bronco.map(lambda row: {'entities': filter_entities(row['entities'], [label])})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3e8d4826-30d4-4752-83fa-fed2ae359fa0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    k1: Dataset({\n",
       "        features: ['id', 'document_id', 'passages', 'entities', 'events', 'coreferences', 'relations'],\n",
       "        num_rows: 1\n",
       "    })\n",
       "    k2: Dataset({\n",
       "        features: ['id', 'document_id', 'passages', 'entities', 'events', 'coreferences', 'relations'],\n",
       "        num_rows: 1\n",
       "    })\n",
       "    k3: Dataset({\n",
       "        features: ['id', 'document_id', 'passages', 'entities', 'events', 'coreferences', 'relations'],\n",
       "        num_rows: 1\n",
       "    })\n",
       "    k4: Dataset({\n",
       "        features: ['id', 'document_id', 'passages', 'entities', 'events', 'coreferences', 'relations'],\n",
       "        num_rows: 1\n",
       "    })\n",
       "    k5: Dataset({\n",
       "        features: ['id', 'document_id', 'passages', 'entities', 'events', 'coreferences', 'relations'],\n",
       "        num_rows: 1\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import DatasetDict\n",
    "\n",
    "ground_truth = DatasetDict()\n",
    "for k in range(5):\n",
    "    ground_truth[f\"k{k+1}\"] = ds[\"train\"].select([k])\n",
    "    \n",
    "ds = ground_truth\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85ec4049-85cb-4a4b-996c-12142d13062a",
   "metadata": {},
   "source": [
    "# Run Candidate Generator\n",
    "We will use the default `EnsembleLinker`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "11662199-832c-49cf-9591-43eb82e5bc75",
   "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\">[10/18/23 19:58:31] </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[10/18/23 19:58:31]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading hierarchical faiss index                                \u001b]8;id=682292;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=417359;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/atc/index/sapbert/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">embed_faiss_h</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\">ier.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=288595;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=41492;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/atc/index/sapbert/\u001b[0m\u001b[95membed_faiss_h\u001b[0m \u001b[2m                   \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[95mier.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\">[10/18/23 19:58:32] </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\">477959</span>                                                        <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/18/23 19:58:32]\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=346283;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=436706;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;36m477959\u001b[0m                                                        \u001b[2m                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from xmen.linkers import default_ensemble\n",
    "from xmen.evaluation import evaluate_at_k\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "index_base_path = Path(f\"{os.path.expanduser('~')}/.cache/xmen/{label2dict[label]}/index\")\n",
    "linker = default_ensemble(index_base_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "12537346-9e8c-4986-8a3b-5f4d068f6294",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.44375\n",
      "Recall@2 0.525\n",
      "Recall@4 0.64375\n",
      "Recall@8 0.7375\n",
      "Recall@16 0.771875\n",
      "Recall@32 0.8125\n",
      "Recall@64 0.81875\n"
     ]
    }
   ],
   "source": [
    "candidates = linker.predict_batch(ds)\n",
    "\n",
    "# Recall for different numbers of candidates (k)\n",
    "_ = evaluate_at_k(ds['k5'], candidates['k5'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7840718b-0581-4267-a809-fe4af2d71bc6",
   "metadata": {},
   "source": [
    "# Train Cross-encoder\n",
    "We use a cross-encoder to embed the mention with their context together with all potential candidates. This way, we can learn the best ranking of candidates from the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "575c872c-9bf7-42a2-aaff-3fa67551a8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.reranking.cross_encoder import CrossEncoderReranker, CrossEncoderTrainingArgs\n",
    "from xmen.data import filter_and_apply_threshold\n",
    "from xmen import load_kb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "32719dc5-35a4-48f9-82f7-d80cab1a4c46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context length: 128\n",
      "Use NIL values: True\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95f8a9a889b34744b8f0641f1a3cc649",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/325 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "760359ddee054ba0b4461bf5c229c1fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/325 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "986038a360e24d92a42d3a6bb469c634",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/325 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "edc6d8affa094a569de39c07c77a2fb9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/333 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6f47051965f40d7bfdbf1292fb0e7d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/333 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0781a291ecb44299da342fbfe858fee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/333 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3f062d598a04d349700086891744597",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/338 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3715863442e14aabb7fe6115fbe21035",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/338 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "818e9edfa43f47b08df70220d14958b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/338 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ac7ef833d7c47ccb00797db476c24cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/314 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba3014f5125b4575b36cd0ba3f3ca6d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/314 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e30e7f3f1719422cad5197417f9cf48b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/314 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "203d1bf971c14f60b5300792fa7b909f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/320 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "feff3c1c337e459282d840d3a66220ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/320 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "842c1462469f4652a7d3a56f60e52026",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/320 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'k1': [325 items],\n",
       " 'k2': [333 items],\n",
       " 'k3': [338 items],\n",
       " 'k4': [314 items],\n",
       " 'k5': [320 items]}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K_RERANKING = 64\n",
    "candidates = filter_and_apply_threshold(candidates, K_RERANKING, 0.0)\n",
    "kb = load_kb(index_base_path.parent / f\"{label2dict[label]}.jsonl\")\n",
    "\n",
    "cross_enc_ds = CrossEncoderReranker.prepare_data(candidates, ds, kb)\n",
    "cross_enc_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daae0c0c-da20-44fd-bed3-10afb8822049",
   "metadata": {},
   "source": [
    "Now we set the training arguments, train-eval splits and fit the model. Depending on the number of epochs, training can take several hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6f959029-0cf1-4e8a-a11d-1f245c5b1347",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Choose train and evaluation folds\n",
    "train_folds = [\"k1\", \"k2\", \"k3\"]\n",
    "val_fold = \"k4\"\n",
    "test_fold = \"k5\"\n",
    "\n",
    "train = sum([cross_enc_ds[k].dataset for k in train_folds],[])\n",
    "val = cross_enc_ds[val_fold].dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "105a6518-ede9-4132-900f-72a0615dab75",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_args = CrossEncoderTrainingArgs(num_train_epochs = 5)\n",
    "\n",
    "rr = CrossEncoderReranker()\n",
    "output_dir = f'../outputs/{label2dict[label]}_index/cross_encoder_training/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c319a6fb-d34e-4df7-a322-c8f179424e3c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model_name := bert-base-multilingual-cased\n",
      "num_train_epochs := 5\n",
      "fp16 := True\n",
      "label_smoothing := False\n",
      "rank_regularization := 1.0\n",
      "train_layers := None\n",
      "softmax_loss := True\n",
      "random_seed := 42\n",
      "learning_rate := 2e-05\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForSequenceClassification: ['cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-multilingual-cased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-10-18 20:19:11 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 0:\n",
      "2023-10-18 20:22:01 - Accuracy: 0.5191082802547771\n",
      "2023-10-18 20:22:01 - Accuracy @ 5: 0.7643312101910829\n",
      "2023-10-18 20:22:01 - Accuracy @ 64: 1.0\n",
      "2023-10-18 20:22:01 - Baseline Accuracy: 0.4394904458598726\n",
      "2023-10-18 20:22:01 - Save model to ../outputs/atc_index/cross_encoder_training/\n",
      "2023-10-18 20:37:41 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 1:\n",
      "2023-10-18 20:40:31 - Accuracy: 0.6019108280254777\n",
      "2023-10-18 20:40:31 - Accuracy @ 5: 0.8757961783439491\n",
      "2023-10-18 20:40:31 - Accuracy @ 64: 1.0\n",
      "2023-10-18 20:40:31 - Baseline Accuracy: 0.4394904458598726\n",
      "2023-10-18 20:40:31 - Save model to ../outputs/atc_index/cross_encoder_training/\n",
      "2023-10-18 20:56:12 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 2:\n",
      "2023-10-18 20:59:02 - Accuracy: 0.60828025477707\n",
      "2023-10-18 20:59:02 - Accuracy @ 5: 0.8949044585987261\n",
      "2023-10-18 20:59:02 - Accuracy @ 64: 1.0\n",
      "2023-10-18 20:59:02 - Baseline Accuracy: 0.4394904458598726\n",
      "2023-10-18 20:59:02 - Save model to ../outputs/atc_index/cross_encoder_training/\n",
      "2023-10-18 21:14:40 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 3:\n",
      "2023-10-18 21:17:31 - Accuracy: 0.6337579617834395\n",
      "2023-10-18 21:17:31 - Accuracy @ 5: 0.8980891719745223\n",
      "2023-10-18 21:17:31 - Accuracy @ 64: 1.0\n",
      "2023-10-18 21:17:31 - Baseline Accuracy: 0.4394904458598726\n",
      "2023-10-18 21:17:31 - Save model to ../outputs/atc_index/cross_encoder_training/\n",
      "2023-10-18 21:33:11 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 4:\n",
      "2023-10-18 21:36:01 - Accuracy: 0.6369426751592356\n",
      "2023-10-18 21:36:01 - Accuracy @ 5: 0.89171974522293\n",
      "2023-10-18 21:36:01 - Accuracy @ 64: 1.0\n",
      "2023-10-18 21:36:01 - Baseline Accuracy: 0.4394904458598726\n",
      "2023-10-18 21:36:01 - Save model to ../outputs/atc_index/cross_encoder_training/\n"
     ]
    }
   ],
   "source": [
    "rr.fit(\n",
    "    train_dataset = train,\n",
    "    val_dataset = val,\n",
    "    output_dir= output_dir,\n",
    "    training_args = train_args,\n",
    "    show_progress_bar = False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e84b2ad-4e5c-4af0-be37-3e29a4cc1a6f",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Evaluate Cross-encoder\n",
    "Now we can take our trained model and test it on data outside of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "056b9c8d-f364-420f-8d54-b3986f8b567c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-10-18 21:36:07 - Use pytorch device: cuda\n"
     ]
    }
   ],
   "source": [
    "rr = CrossEncoderReranker.load(output_dir, device=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d3507f0e-8d45-4d78-a662-75a60363e19d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12424b1afdf64c2ba3e08f9d61a2089e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/314 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.7197452229299363\n",
      "Recall@2 0.7547770700636943\n",
      "Recall@4 0.7770700636942676\n",
      "Recall@8 0.7866242038216561\n",
      "Recall@16 0.7993630573248408\n",
      "Recall@32 0.802547770700637\n",
      "Recall@64 0.802547770700637\n"
     ]
    }
   ],
   "source": [
    "cross_enc_pred_val = rr.rerank_batch(candidates[val_fold], cross_enc_ds[val_fold])\n",
    "_ = evaluate_at_k(ds[val_fold], cross_enc_pred_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f9f97c16-f19c-4e54-9b4d-108a3c8f83c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b8b124e1554c4baeabc9ab4ba469633a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/320 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.675\n",
      "Recall@2 0.725\n",
      "Recall@4 0.734375\n",
      "Recall@8 0.78125\n",
      "Recall@16 0.790625\n",
      "Recall@32 0.8\n",
      "Recall@64 0.80625\n"
     ]
    }
   ],
   "source": [
    "cross_enc_pred_test = rr.rerank_batch(candidates[test_fold], cross_enc_ds[test_fold])\n",
    "_ = evaluate_at_k(ds[test_fold], cross_enc_pred_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa344887-7a73-4697-a434-12c43aef97b3",
   "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
}