notebooks/00_Getting_Started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "85c0c045-2900-47d5-8c9b-5ebe0b4c01c1",
   "metadata": {},
   "source": [
    "# Getting Started\n",
    "\n",
    "In this notebook, we will walk you through a minimal example for candidate generation and re-reranking with the [Quaero](https://quaerofrenchmed.limsi.fr/) corpus."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a05ddc-a8b0-444d-890b-898ce9e5f302",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "Go to the root folder of the repository and run the following commands\n",
    "\n",
    "`xmen dict examples/conf/quaero.yaml`\n",
    "\n",
    "This command creates the Quaero-specific UMLS subset as a jsonl file.\n",
    "If you don't have access to UMLS metathesaurus 2014AB release, you can adapt the yaml file to point to any other release.\n",
    "\n",
    "`xmen index examples/conf/quaero.yaml --all`\n",
    "\n",
    "This command creates the indices used candidate generations (both TF-IDF and SapBERT).\n",
    "\n",
    "By default, all files are written to `~/.cache/xmen`, unless configured otherwise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "482b215a-9a15-4227-8c88-da811c0c0fe9",
   "metadata": {},
   "source": [
    "## Loading Knowledge Base and Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "46917412-c6e3-48c6-8e45-278aab09202c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "base_path = Path.home() / \".cache/xmen/quaero\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8c419ff7-306f-4522-811c-cffadea0010d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset quaero (/home/Florian.Borchert/.cache/huggingface/datasets/bigbio___quaero/quaero_medline_bigbio_kb/1.0.0/c215a01af5fd15163e392e2f78a4f67f1fe3567d92f53b7f907c95e55dcdfe93)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "faa23a0a4d494b0aaefc5b2f7695ec36",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load Quaero MEDLINE subset through BigBIO dataloader\n",
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"bigbio/quaero\", \"quaero_medline_bigbio_kb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "540bb862-7d66-4216-bf22-1943b2fa654c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load knowledge base\n",
    "from xmen import load_kb\n",
    "kb = load_kb(base_path / \"quaero.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad987400-dfc7-4ecc-8fd0-5aeb408904e2",
   "metadata": {},
   "source": [
    "# Candidate Generation\n",
    "\n",
    "We use the pre-computed indices to retrieve the 64 most similar concepts for each mention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "37d5d24c-7bbf-4254-9d74-52b983c24d12",
   "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 12:38: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 12:38:31]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO    \u001b[0m Loading hierarchical faiss index                                \u001b]8;id=392403;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=659557;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/quaero/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=729210;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=552471;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/quaero/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\">[10/18/23 12:39:06] </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\">6910830</span>                                                       <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[10/18/23 12:39:06]\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=461408;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=307951;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;36m6910830\u001b[0m                                                       \u001b[2m                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Generate candidates with Ensemble of TF-IDF + SapBERT\n",
    "from xmen.linkers import default_ensemble\n",
    "linker = default_ensemble(base_path / \"index\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "276af3a6-e245-4364-9ef6-d9fd8cbb862a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/833 [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/833 [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/832 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "candidates = linker.predict_batch(dataset, top_k=64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4526ee50-a9d1-48c1-8e08-f676f2ef6da4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'strict': {'precision': 0.5802389409105586,\n",
       "  'recall': 0.5804263565891473,\n",
       "  'fscore': 0.5803326336186015,\n",
       "  'ptp': 1797,\n",
       "  'fp': 1300,\n",
       "  'rtp': 1797,\n",
       "  'fn': 1299,\n",
       "  'n_docs_system': 833,\n",
       "  'n_annos_system': 3103,\n",
       "  'n_docs_gold': 833,\n",
       "  'n_annos_gold': 3097}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from xmen.evaluation import evaluate\n",
    "evaluate(dataset['test'], candidates['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4da2e1d1-2ed2-4f14-b3f7-a19a2e45190f",
   "metadata": {},
   "source": [
    "# Semantic Group Filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c4249662",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/833 [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/833 [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/832 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from xmen.data import SemanticGroupFilter\n",
    "candidates = SemanticGroupFilter(kb).transform_batch(candidates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "86221f29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'strict': {'precision': 0.6606276285991588,\n",
       "  'recall': 0.6595607235142119,\n",
       "  'fscore': 0.6600937449490868,\n",
       "  'ptp': 2042,\n",
       "  'fp': 1049,\n",
       "  'rtp': 2042,\n",
       "  'fn': 1054,\n",
       "  'n_docs_system': 833,\n",
       "  'n_annos_system': 3095,\n",
       "  'n_docs_gold': 833,\n",
       "  'n_annos_gold': 3097}}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluate(dataset['test'], candidates['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8617e85-44af-456d-a9eb-62e2ea6d9a88",
   "metadata": {},
   "source": [
    "# Reranking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f26c4831-7b61-4e68-8bd6-2c40c8ecfe6d",
   "metadata": {},
   "source": [
    "## Using a Pre-trained Model for Reranking\n",
    "\n",
    "When no are little annotated training data is available, a pre-trained model can provide good re-ranking performance.\n",
    "Here, we use a cross-encoder that was trained on an automatically translated version of MedMentions (see the notebook [01_Translation.ipynb](01_Translation.ipynb))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cacadea0-079d-4d04-8ca0-d56de356c431",
   "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": "9b66a4ca119e4553ab367ebe81b6922e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2994 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e9d155c31624dd88202908d0db9336f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2994 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7cb6d8432fe495881c3b5bb8920afc0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2994 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1fca1ee972704559a7cacec42968593a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3103 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1aafb283ccd4ba98eb6a89bbbc8f022",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3103 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22b9cfbd9197487a92af68d3cd5acacc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3103 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d1db7fd6767f4407ae77707a23dc81f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2977 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de246595d65f49e88dc09f84bb028656",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2977 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d14d85dd4a4741adba814804b5255a63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2977 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from xmen.reranking.cross_encoder import *\n",
    "ce_dataset = CrossEncoderReranker.prepare_data(candidates, dataset, kb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "82de6969-0899-4422-b06c-7f95db789299",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f86ca658546f4ba2bed562361084b828",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/731 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "651fc9a8876c44b6bf29b33c8b233621",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/438M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9635b63c4ee54ba1a68622daf834decc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/494 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38988beeeb3349599b180e2b38f80496",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)solve/main/vocab.txt:   0%|          | 0.00/226k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e2249535247f4deeb0568aa737a7c2de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/707k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f15cdd48c8864889b4eb9f16cc6916f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)in/added_tokens.json:   0%|          | 0.00/85.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0e0644a01ed4592b5787bc6160fbcb3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/193 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "# CE re-ranker, pre-trained on French translation of MedMentions\n",
    "rr = CrossEncoderReranker.load(\"phlobo/xmen-fr-ce-medmentions\", device=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "21d793c8-6b05-47b4-8f3f-6948fb46c15f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a62dc437b99c4c2797dfc10c4071dea5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/3095 [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/833 [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/833 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Predict on test set\n",
    "prediction = rr.rerank_batch(candidates[\"test\"], ce_dataset[\"test\"], allow_nil=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a818b24d-c3b9-4290-8dbc-7f0d74bbcf75",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'strict': {'precision': 0.7234248788368336,\n",
       "  'recall': 0.7231912144702842,\n",
       "  'fscore': 0.7233080277822647,\n",
       "  'ptp': 2239,\n",
       "  'fp': 856,\n",
       "  'rtp': 2239,\n",
       "  'fn': 857,\n",
       "  'n_docs_system': 833,\n",
       "  'n_annos_system': 3095,\n",
       "  'n_docs_gold': 833,\n",
       "  'n_annos_gold': 3097}}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluate(dataset['test'], prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74aef74d-58a7-400e-ba8f-a31e4a9fb3b6",
   "metadata": {},
   "source": [
    "## Training a Fully-supervised Model\n",
    "\n",
    "Finally, with enough training data, we can also train a fully supervsied cross-encoder to learn a better task-specific ranking.\n",
    "\n",
    "The cross-encoder is trained with batches of 64 candidates. \n",
    "We keep the checkpoint that maximizes accuracy@1 on the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ab585f3c-c829-4fbb-912f-c826545a362b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model_name := bert-base-multilingual-cased\n",
      "num_train_epochs := 3\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 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 13:49:28 - Use pytorch device: cuda\n",
      "2023-10-18 14:08:02 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 0:\n",
      "2023-10-18 14:18:57 - Accuracy: 0.7233468286099866\n",
      "2023-10-18 14:18:57 - Accuracy @ 5: 0.861336032388664\n",
      "2023-10-18 14:18:57 - Accuracy @ 64: 0.9946018893387314\n",
      "2023-10-18 14:18:57 - Baseline Accuracy: 0.6686909581646424\n",
      "2023-10-18 14:18:57 - Save model to ./output/cross_encoder\n",
      "2023-10-18 14:37:28 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 1:\n",
      "2023-10-18 14:48:25 - Accuracy: 0.7435897435897436\n",
      "2023-10-18 14:48:25 - Accuracy @ 5: 0.8711201079622132\n",
      "2023-10-18 14:48:25 - Accuracy @ 64: 0.9946018893387314\n",
      "2023-10-18 14:48:25 - Baseline Accuracy: 0.6686909581646424\n",
      "2023-10-18 14:48:25 - Save model to ./output/cross_encoder\n",
      "2023-10-18 15:07:04 - EntityLinkingEvaluator: Evaluating the model on eval dataset after epoch 2:\n",
      "2023-10-18 15:18:02 - Accuracy: 0.7510121457489879\n",
      "2023-10-18 15:18:02 - Accuracy @ 5: 0.8734817813765182\n",
      "2023-10-18 15:18:02 - Accuracy @ 64: 0.9946018893387314\n",
      "2023-10-18 15:18:02 - Baseline Accuracy: 0.6686909581646424\n",
      "2023-10-18 15:18:02 - Save model to ./output/cross_encoder\n"
     ]
    }
   ],
   "source": [
    "# Train the cross-encoder on the training set and use the validation set for model selection\n",
    "rr_fs = CrossEncoderReranker()\n",
    "# Note: we train for 3 epochs only (you likely want to increase this for better performance)\n",
    "args = CrossEncoderTrainingArgs(num_train_epochs=3)\n",
    "rr_fs.fit(args, ce_dataset[\"train\"].dataset, ce_dataset[\"validation\"].dataset, show_progress_bar=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d09c3b36-bf23-4284-a67f-d832614f2424",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89179d97be4843c5a2254fbc82ad8f47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/3095 [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/833 [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/833 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Predict on test set\n",
    "prediction_fs = rr_fs.rerank_batch(candidates[\"test\"], ce_dataset[\"test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8a75f0-7b64-4cdb-87f3-f6b739cb320e",
   "metadata": {},
   "source": [
    "# Evaluation\n",
    "\n",
    "## Metrics\n",
    "\n",
    "We compute metrics for the predicted concepts vs. the gold-standard for different values of k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8ac9da12-8d38-4364-96f2-be5a1e1b9fc5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'strict': {'precision': 0.7528015820698748,\n",
       "  'recall': 0.7377260981912145,\n",
       "  'fscore': 0.7451876019575857,\n",
       "  'ptp': 2284,\n",
       "  'fp': 750,\n",
       "  'rtp': 2284,\n",
       "  'fn': 812,\n",
       "  'n_docs_system': 833,\n",
       "  'n_annos_system': 3034,\n",
       "  'n_docs_gold': 833,\n",
       "  'n_annos_gold': 3097}}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluate(dataset['test'], prediction_fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f38d2530-642f-4d0a-95ec-94f21de52483",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.6595607235142119\n",
      "Recall@2 0.790374677002584\n",
      "Recall@4 0.8449612403100775\n",
      "Recall@8 0.8656330749354005\n",
      "Recall@16 0.8853359173126615\n",
      "Recall@32 0.8963178294573644\n",
      "Recall@64 0.8998708010335917\n"
     ]
    }
   ],
   "source": [
    "from xmen.evaluation import evaluate_at_k\n",
    "_ = evaluate_at_k(dataset['test'], candidates['test'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "12421a31-eb88-4394-85d1-42295178ed4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@1 0.7377260981912145\n",
      "Recall@2 0.8049095607235142\n",
      "Recall@4 0.8397932816537468\n",
      "Recall@8 0.8601421188630491\n",
      "Recall@16 0.8782299741602068\n",
      "Recall@32 0.8888888888888888\n",
      "Recall@64 0.8934108527131783\n"
     ]
    }
   ],
   "source": [
    "_ = evaluate_at_k(dataset['test'], prediction_fs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3f94940-fb36-4d72-adf5-6e3e60450a19",
   "metadata": {},
   "source": [
    "## Error Analysis\n",
    "\n",
    "To gain more insights into the ranking errors (and potential NER errors, which do not apply here), we use the `error_analysis` functionality of xmen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b84234f7-ee2c-449d-988c-6c6b60f2dd5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xmen.evaluation import error_analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fc535571-0e7c-49f6-a87b-6d61aa0c3c2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>_word_len</th>\n",
       "      <th>_abbrev</th>\n",
       "      <th>gt_start</th>\n",
       "      <th>gt_end</th>\n",
       "      <th>gt_text</th>\n",
       "      <th>gold_concept</th>\n",
       "      <th>gold_type</th>\n",
       "      <th>pred_index</th>\n",
       "      <th>pred_index_score</th>\n",
       "      <th>pred_top</th>\n",
       "      <th>pred_top_score</th>\n",
       "      <th>document_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>[Etude]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0008972'}</td>\n",
       "      <td>PROC</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C0013621</td>\n",
       "      <td>0.692185</td>\n",
       "      <td>1009657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>47</td>\n",
       "      <td>55</td>\n",
       "      <td>[Africain]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0027567'}</td>\n",
       "      <td>LIVB</td>\n",
       "      <td>0</td>\n",
       "      <td>0.863724</td>\n",
       "      <td>C0027567</td>\n",
       "      <td>0.863724</td>\n",
       "      <td>1009657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>4</td>\n",
       "      <td>13</td>\n",
       "      <td>[gangrènes]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0017086'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>0.973051</td>\n",
       "      <td>C0017086</td>\n",
       "      <td>0.973051</td>\n",
       "      <td>1009808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>4</td>\n",
       "      <td>22</td>\n",
       "      <td>[gangrènes gazeuses]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0017105'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>0.935468</td>\n",
       "      <td>C0017105</td>\n",
       "      <td>0.935468</td>\n",
       "      <td>1009808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>14</td>\n",
       "      <td>22</td>\n",
       "      <td>[gazeuses]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0017110'}</td>\n",
       "      <td>CHEM</td>\n",
       "      <td>2</td>\n",
       "      <td>0.808053</td>\n",
       "      <td>C0537495</td>\n",
       "      <td>0.858877</td>\n",
       "      <td>1009808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>35</td>\n",
       "      <td>46</td>\n",
       "      <td>[tuberculose]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0041296'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>C0041296</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>35</td>\n",
       "      <td>53</td>\n",
       "      <td>[tuberculose bovine]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0041307'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>C0041307</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>47</td>\n",
       "      <td>53</td>\n",
       "      <td>[bovine]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0007452'}</td>\n",
       "      <td>LIVB</td>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>C3667982</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>56</td>\n",
       "      <td>75</td>\n",
       "      <td>[Mycobacterium bovis]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0026917'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C0041307</td>\n",
       "      <td>0.825099</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>78</td>\n",
       "      <td>86</td>\n",
       "      <td>[observés]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0302523'}</td>\n",
       "      <td>PROC</td>\n",
       "      <td>1</td>\n",
       "      <td>0.852574</td>\n",
       "      <td>C0700325</td>\n",
       "      <td>0.852574</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   _word_len  _abbrev gt_start gt_end                gt_text  \\\n",
       "0          1    False        0      5                [Etude]   \n",
       "1          1    False       47     55             [Africain]   \n",
       "2          1    False        4     13            [gangrènes]   \n",
       "3          2    False        4     22   [gangrènes gazeuses]   \n",
       "4          1    False       14     22             [gazeuses]   \n",
       "5          1    False       35     46          [tuberculose]   \n",
       "6          2    False       35     53   [tuberculose bovine]   \n",
       "7          1    False       47     53               [bovine]   \n",
       "8          2    False       56     75  [Mycobacterium bovis]   \n",
       "9          1    False       78     86             [observés]   \n",
       "\n",
       "                               gold_concept gold_type pred_index  \\\n",
       "0  {'db_name': 'UMLS', 'db_id': 'C0008972'}      PROC         -1   \n",
       "1  {'db_name': 'UMLS', 'db_id': 'C0027567'}      LIVB          0   \n",
       "2  {'db_name': 'UMLS', 'db_id': 'C0017086'}      DISO          0   \n",
       "3  {'db_name': 'UMLS', 'db_id': 'C0017105'}      DISO          0   \n",
       "4  {'db_name': 'UMLS', 'db_id': 'C0017110'}      CHEM          2   \n",
       "5  {'db_name': 'UMLS', 'db_id': 'C0041296'}      DISO          0   \n",
       "6  {'db_name': 'UMLS', 'db_id': 'C0041307'}      DISO          0   \n",
       "7  {'db_name': 'UMLS', 'db_id': 'C0007452'}      LIVB          3   \n",
       "8  {'db_name': 'UMLS', 'db_id': 'C0026917'}      DISO         -1   \n",
       "9  {'db_name': 'UMLS', 'db_id': 'C0302523'}      PROC          1   \n",
       "\n",
       "  pred_index_score  pred_top pred_top_score document_id  \n",
       "0              NaN  C0013621       0.692185     1009657  \n",
       "1         0.863724  C0027567       0.863724     1009657  \n",
       "2         0.973051  C0017086       0.973051     1009808  \n",
       "3         0.935468  C0017105       0.935468     1009808  \n",
       "4         0.808053  C0537495       0.858877     1009808  \n",
       "5              1.0  C0041296            1.0     1019393  \n",
       "6              1.0  C0041307            1.0     1019393  \n",
       "7              1.0  C3667982            1.0     1019393  \n",
       "8              NaN  C0041307       0.825099     1019393  \n",
       "9         0.852574  C0700325       0.852574     1019393  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Before re-ranking\n",
    "ea_df = error_analysis(dataset['test'], candidates['test'], tasks=['nen'])\n",
    "ea_df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "56e3d2a6-0cb5-43d6-94e5-e5a25c36926b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>_word_len</th>\n",
       "      <th>_abbrev</th>\n",
       "      <th>gt_start</th>\n",
       "      <th>gt_end</th>\n",
       "      <th>gt_text</th>\n",
       "      <th>gold_concept</th>\n",
       "      <th>gold_type</th>\n",
       "      <th>pred_index</th>\n",
       "      <th>pred_index_score</th>\n",
       "      <th>pred_top</th>\n",
       "      <th>pred_top_score</th>\n",
       "      <th>document_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>[Etude]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0008972'}</td>\n",
       "      <td>PROC</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NIL</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1009657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>47</td>\n",
       "      <td>55</td>\n",
       "      <td>[Africain]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0027567'}</td>\n",
       "      <td>LIVB</td>\n",
       "      <td>1</td>\n",
       "      <td>0.037805</td>\n",
       "      <td>C0085756</td>\n",
       "      <td>0.041022</td>\n",
       "      <td>1009657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>4</td>\n",
       "      <td>13</td>\n",
       "      <td>[gangrènes]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0017086'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>0.042884</td>\n",
       "      <td>C0017086</td>\n",
       "      <td>0.042884</td>\n",
       "      <td>1009808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>4</td>\n",
       "      <td>22</td>\n",
       "      <td>[gangrènes gazeuses]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0017105'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>0.061593</td>\n",
       "      <td>C0017105</td>\n",
       "      <td>0.061593</td>\n",
       "      <td>1009808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>14</td>\n",
       "      <td>22</td>\n",
       "      <td>[gazeuses]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0017110'}</td>\n",
       "      <td>CHEM</td>\n",
       "      <td>0</td>\n",
       "      <td>0.090157</td>\n",
       "      <td>C0017110</td>\n",
       "      <td>0.090157</td>\n",
       "      <td>1009808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>35</td>\n",
       "      <td>46</td>\n",
       "      <td>[tuberculose]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0041296'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>0.04742</td>\n",
       "      <td>C0041296</td>\n",
       "      <td>0.04742</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>35</td>\n",
       "      <td>53</td>\n",
       "      <td>[tuberculose bovine]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0041307'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>0</td>\n",
       "      <td>0.062842</td>\n",
       "      <td>C0041307</td>\n",
       "      <td>0.062842</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>47</td>\n",
       "      <td>53</td>\n",
       "      <td>[bovine]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0007452'}</td>\n",
       "      <td>LIVB</td>\n",
       "      <td>0</td>\n",
       "      <td>0.087439</td>\n",
       "      <td>C0007452</td>\n",
       "      <td>0.087439</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2</td>\n",
       "      <td>False</td>\n",
       "      <td>56</td>\n",
       "      <td>75</td>\n",
       "      <td>[Mycobacterium bovis]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0026917'}</td>\n",
       "      <td>DISO</td>\n",
       "      <td>-1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C0041307</td>\n",
       "      <td>0.550823</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1</td>\n",
       "      <td>False</td>\n",
       "      <td>78</td>\n",
       "      <td>86</td>\n",
       "      <td>[observés]</td>\n",
       "      <td>{'db_name': 'UMLS', 'db_id': 'C0302523'}</td>\n",
       "      <td>PROC</td>\n",
       "      <td>0</td>\n",
       "      <td>0.050934</td>\n",
       "      <td>C0302523</td>\n",
       "      <td>0.050934</td>\n",
       "      <td>1019393</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   _word_len  _abbrev gt_start gt_end                gt_text  \\\n",
       "0          1    False        0      5                [Etude]   \n",
       "1          1    False       47     55             [Africain]   \n",
       "2          1    False        4     13            [gangrènes]   \n",
       "3          2    False        4     22   [gangrènes gazeuses]   \n",
       "4          1    False       14     22             [gazeuses]   \n",
       "5          1    False       35     46          [tuberculose]   \n",
       "6          2    False       35     53   [tuberculose bovine]   \n",
       "7          1    False       47     53               [bovine]   \n",
       "8          2    False       56     75  [Mycobacterium bovis]   \n",
       "9          1    False       78     86             [observés]   \n",
       "\n",
       "                               gold_concept gold_type pred_index  \\\n",
       "0  {'db_name': 'UMLS', 'db_id': 'C0008972'}      PROC         -1   \n",
       "1  {'db_name': 'UMLS', 'db_id': 'C0027567'}      LIVB          1   \n",
       "2  {'db_name': 'UMLS', 'db_id': 'C0017086'}      DISO          0   \n",
       "3  {'db_name': 'UMLS', 'db_id': 'C0017105'}      DISO          0   \n",
       "4  {'db_name': 'UMLS', 'db_id': 'C0017110'}      CHEM          0   \n",
       "5  {'db_name': 'UMLS', 'db_id': 'C0041296'}      DISO          0   \n",
       "6  {'db_name': 'UMLS', 'db_id': 'C0041307'}      DISO          0   \n",
       "7  {'db_name': 'UMLS', 'db_id': 'C0007452'}      LIVB          0   \n",
       "8  {'db_name': 'UMLS', 'db_id': 'C0026917'}      DISO         -1   \n",
       "9  {'db_name': 'UMLS', 'db_id': 'C0302523'}      PROC          0   \n",
       "\n",
       "  pred_index_score  pred_top pred_top_score document_id  \n",
       "0              NaN       NIL            NaN     1009657  \n",
       "1         0.037805  C0085756       0.041022     1009657  \n",
       "2         0.042884  C0017086       0.042884     1009808  \n",
       "3         0.061593  C0017105       0.061593     1009808  \n",
       "4         0.090157  C0017110       0.090157     1009808  \n",
       "5          0.04742  C0041296        0.04742     1019393  \n",
       "6         0.062842  C0041307       0.062842     1019393  \n",
       "7         0.087439  C0007452       0.087439     1019393  \n",
       "8              NaN  C0041307       0.550823     1019393  \n",
       "9         0.050934  C0302523       0.050934     1019393  "
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# After Re-Ranking\n",
    "ea_df = error_analysis(dataset['test'], prediction_fs, tasks=['nen'])\n",
    "ea_df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d473f20a-998e-401b-bc16-42466d4de8b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: xlabel='pred_index'>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjAAAAGwCAYAAAC3qV8qAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAmfUlEQVR4nO3deXTU9b3/8ddkD8JMCJDtNiQglQQriAFDlFVSAqSoBY+CKHovJZWTKDbK1nrD1spSK9JiodpqrEWFtrjBMSXGC+nFABJvXGKMQcHghQQKTYYgez6/P/wxlylbAhknn/B8nPM9h5nvZ+b7/prSPJmZb+IwxhgBAABYJMDfAwAAADQXAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6wT5ewBfaWxs1J49e9ShQwc5HA5/jwMAAJrAGKNDhw4pLi5OAQHnf52lzQbMnj17FB8f7+8xAADAJdi9e7e+853vnHd/mw2YDh06SPrmP4DT6fTzNAAAoCncbrfi4+M938fPp80GzOm3jZxOJwEDAIBlLvbxDz7ECwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOkH+HgAtL3HWep8fY9eiTJ8fAwCA8+EVGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1mlWwCxcuFD9+/dXhw4dFBUVpdtvv12VlZVea44ePars7Gx16tRJ7du317hx41RbW+u1prq6WpmZmWrXrp2ioqI0ffp0nTx50mvNxo0bdcMNNyg0NFQ9evRQfn7+pZ0hAABoc5oVMJs2bVJ2dra2bNmiwsJCnThxQiNGjNDhw4c9a37yk5/ozTff1J///Gdt2rRJe/bs0dixYz37T506pczMTB0/flzvvvuuXnjhBeXn5ysvL8+zZufOncrMzNSwYcNUVlamhx9+WD/60Y/0t7/9rQVOGQAA2M5hjDGX+uD9+/crKipKmzZt0uDBg1VfX68uXbropZde0h133CFJ+vTTT5WcnKySkhINGDBAb731ln7wgx9oz549io6OliStXLlSM2fO1P79+xUSEqKZM2dq/fr1+vjjjz3HGj9+vOrq6lRQUNCk2dxut1wul+rr6+V0Oi/1FK2UOGu9z4+xa1Gmz48BALjyNPX792V9Bqa+vl6SFBkZKUkqLS3ViRMnlJ6e7lmTlJSkrl27qqSkRJJUUlKi6667zhMvkpSRkSG3263y8nLPmjOf4/Sa089xLseOHZPb7fbaAABA23TJAdPY2KiHH35YN998s773ve9JkmpqahQSEqKIiAivtdHR0aqpqfGsOTNeTu8/ve9Ca9xut44cOXLOeRYuXCiXy+XZ4uPjL/XUAABAK3fJAZOdna2PP/5Yr7zySkvOc8lmz56t+vp6z7Z7925/jwQAAHwk6FIelJOTo3Xr1qm4uFjf+c53PPfHxMTo+PHjqqur83oVpra2VjExMZ4127Zt83q+01cpnbnmX69cqq2tldPpVHh4+DlnCg0NVWho6KWcDgAAsEyzXoExxignJ0evvvqq3nnnHXXr1s1rf0pKioKDg1VUVOS5r7KyUtXV1UpLS5MkpaWl6aOPPtK+ffs8awoLC+V0OtWrVy/PmjOf4/Sa088BAACubM16BSY7O1svvfSSXn/9dXXo0MHzmRWXy6Xw8HC5XC5NnjxZubm5ioyMlNPp1IMPPqi0tDQNGDBAkjRixAj16tVL9957r5YsWaKamho99thjys7O9ryC8sADD2j58uWaMWOG/uM//kPvvPOO1qxZo/XrfX91DQAAaP2a9QrMihUrVF9fr6FDhyo2NtazrV692rNm6dKl+sEPfqBx48Zp8ODBiomJ0dq1az37AwMDtW7dOgUGBiotLU333HOPJk2apPnz53vWdOvWTevXr1dhYaH69OmjX/3qV/r973+vjIyMFjhlAABgu8v6OTCtGT8Hxrf4OTAAAF/4Vn4ODAAAgD8QMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOs0OmOLiYo0ZM0ZxcXFyOBx67bXXvPbff//9cjgcXtvIkSO91hw8eFATJ06U0+lURESEJk+erIaGBq81H374oQYNGqSwsDDFx8dryZIlzT87AADQJjU7YA4fPqw+ffro6aefPu+akSNHau/evZ7t5Zdf9to/ceJElZeXq7CwUOvWrVNxcbGysrI8+91ut0aMGKGEhASVlpbql7/8pebOnatnnnmmueMCAIA2KKi5Dxg1apRGjRp1wTWhoaGKiYk5576KigoVFBTovffeU79+/SRJv/nNbzR69Gg98cQTiouL06pVq3T8+HE999xzCgkJ0bXXXquysjI9+eSTXqEDAACuTD75DMzGjRsVFRWlnj17aurUqTpw4IBnX0lJiSIiIjzxIknp6ekKCAjQ1q1bPWsGDx6skJAQz5qMjAxVVlbqn//85zmPeezYMbndbq8NAAC0TS0eMCNHjtQf//hHFRUVafHixdq0aZNGjRqlU6dOSZJqamoUFRXl9ZigoCBFRkaqpqbGsyY6Otprzenbp9f8q4ULF8rlcnm2+Pj4lj41AADQSjT7LaSLGT9+vOfP1113nXr37q2rr75aGzdu1PDhw1v6cB6zZ89Wbm6u57bb7SZiAABoo3x+GXX37t3VuXNn7dixQ5IUExOjffv2ea05efKkDh486PncTExMjGpra73WnL59vs/WhIaGyul0em0AAKBt8nnAfPXVVzpw4IBiY2MlSWlpaaqrq1NpaalnzTvvvKPGxkalpqZ61hQXF+vEiROeNYWFherZs6c6duzo65EBAEAr1+yAaWhoUFlZmcrKyiRJO3fuVFlZmaqrq9XQ0KDp06dry5Yt2rVrl4qKinTbbbepR48eysjIkCQlJydr5MiRmjJlirZt26bNmzcrJydH48ePV1xcnCTp7rvvVkhIiCZPnqzy8nKtXr1ay5Yt83qLCAAAXLmaHTDbt29X37591bdvX0lSbm6u+vbtq7y8PAUGBurDDz/UrbfeqmuuuUaTJ09WSkqK/v73vys0NNTzHKtWrVJSUpKGDx+u0aNHa+DAgV4/48XlcmnDhg3auXOnUlJS9MgjjygvL49LqAEAgCTJYYwx/h7CF9xut1wul+rr66+4z8Mkzlrv82PsWpTp82MAAK48Tf3+ze9CAgAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1mh0wxcXFGjNmjOLi4uRwOPTaa6957TfGKC8vT7GxsQoPD1d6erqqqqq81hw8eFATJ06U0+lURESEJk+erIaGBq81H374oQYNGqSwsDDFx8dryZIlzT87AADQJjU7YA4fPqw+ffro6aefPuf+JUuW6Ne//rVWrlyprVu36qqrrlJGRoaOHj3qWTNx4kSVl5ersLBQ69atU3FxsbKysjz73W63RowYoYSEBJWWluqXv/yl5s6dq2eeeeYSThEAALQ1DmOMueQHOxx69dVXdfvtt0v65tWXuLg4PfLII3r00UclSfX19YqOjlZ+fr7Gjx+viooK9erVS++995769esnSSooKNDo0aP11VdfKS4uTitWrNDPfvYz1dTUKCQkRJI0a9Ysvfbaa/r000+bNJvb7ZbL5VJ9fb2cTuelnqKVEmet9/kxdi3K9PkxAABXnqZ+/27Rz8Ds3LlTNTU1Sk9P99zncrmUmpqqkpISSVJJSYkiIiI88SJJ6enpCggI0NatWz1rBg8e7IkXScrIyFBlZaX++c9/nvPYx44dk9vt9toAAEDb1KIBU1NTI0mKjo72uj86Otqzr6amRlFRUV77g4KCFBkZ6bXmXM9x5jH+1cKFC+VyuTxbfHz85Z8QAABoldrMVUizZ89WfX29Z9u9e7e/RwIAAD7SogETExMjSaqtrfW6v7a21rMvJiZG+/bt89p/8uRJHTx40GvNuZ7jzGP8q9DQUDmdTq8NAAC0TS0aMN26dVNMTIyKioo897ndbm3dulVpaWmSpLS0NNXV1am0tNSz5p133lFjY6NSU1M9a4qLi3XixAnPmsLCQvXs2VMdO3ZsyZEBAICFmh0wDQ0NKisrU1lZmaRvPrhbVlam6upqORwOPfzww/r5z3+uN954Qx999JEmTZqkuLg4z5VKycnJGjlypKZMmaJt27Zp8+bNysnJ0fjx4xUXFydJuvvuuxUSEqLJkyervLxcq1ev1rJly5Sbm9tiJw4AAOwV1NwHbN++XcOGDfPcPh0V9913n/Lz8zVjxgwdPnxYWVlZqqur08CBA1VQUKCwsDDPY1atWqWcnBwNHz5cAQEBGjdunH7961979rtcLm3YsEHZ2dlKSUlR586dlZeX5/WzYgAAwJXrsn4OTGvGz4HxLX4ODADAF/zyc2AAAAC+DQQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALBOiwfM3Llz5XA4vLakpCTP/qNHjyo7O1udOnVS+/btNW7cONXW1no9R3V1tTIzM9WuXTtFRUVp+vTpOnnyZEuPCgAALBXkiye99tpr9fbbb//fQYL+7zA/+clPtH79ev35z3+Wy+VSTk6Oxo4dq82bN0uSTp06pczMTMXExOjdd9/V3r17NWnSJAUHB+vxxx/3xbgAAMAyPgmYoKAgxcTEnHV/fX29/vCHP+ill17SLbfcIkl6/vnnlZycrC1btmjAgAHasGGDPvnkE7399tuKjo7W9ddfrwULFmjmzJmaO3euQkJCfDEyAACwiE8+A1NVVaW4uDh1795dEydOVHV1tSSptLRUJ06cUHp6umdtUlKSunbtqpKSEklSSUmJrrvuOkVHR3vWZGRkyO12q7y8/LzHPHbsmNxut9cGAADaphYPmNTUVOXn56ugoEArVqzQzp07NWjQIB06dEg1NTUKCQlRRESE12Oio6NVU1MjSaqpqfGKl9P7T+87n4ULF8rlcnm2+Pj4lj0xAADQarT4W0ijRo3y/Ll3795KTU1VQkKC1qxZo/Dw8JY+nMfs2bOVm5vrue12u4kYAADaKJ9fRh0REaFrrrlGO3bsUExMjI4fP666ujqvNbW1tZ7PzMTExJx1VdLp2+f6XM1poaGhcjqdXhsAAGibfB4wDQ0N+vzzzxUbG6uUlBQFBwerqKjIs7+yslLV1dVKS0uTJKWlpemjjz7Svn37PGsKCwvldDrVq1cvX48LAAAs0OJvIT366KMaM2aMEhIStGfPHs2ZM0eBgYGaMGGCXC6XJk+erNzcXEVGRsrpdOrBBx9UWlqaBgwYIEkaMWKEevXqpXvvvVdLlixRTU2NHnvsMWVnZys0NLSlxwUAABZq8YD56quvNGHCBB04cEBdunTRwIEDtWXLFnXp0kWStHTpUgUEBGjcuHE6duyYMjIy9Nvf/tbz+MDAQK1bt05Tp05VWlqarrrqKt13332aP39+S48KAAAs5TDGGH8P4Qtut1sul0v19fVX3OdhEmet9/kxdi3K9PkxAABXnqZ+/+Z3IQEAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOgQMAACwDgEDAACsQ8AAAADrEDAAAMA6BAwAALAOAQMAAKxDwAAAAOsQMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAADAOkH+HqA1SZy13ufH2LUo0+fHAACgreMVGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWCfL3AMC5JM5a7/Nj7FqU6fNjAAB8g1dgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYB0CBgAAWIeAAQAA1iFgAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFiHgAEAANYhYAAAgHUIGAAAYJ0gfw8AtGWJs9b7/Bi7FmX6/BgA0NoQMAAuytchRoQBaC7eQgIAANYhYAAAgHUIGAAAYJ1WHTBPP/20EhMTFRYWptTUVG3bts3fIwEAgFag1X6Id/Xq1crNzdXKlSuVmpqqp556ShkZGaqsrFRUVJS/xwNgGa4Iaz34WqAltNqAefLJJzVlyhT9+7//uyRp5cqVWr9+vZ577jnNmjXLz9MBgH/wzb/1aCtfC1uvMmyVAXP8+HGVlpZq9uzZnvsCAgKUnp6ukpKScz7m2LFjOnbsmOd2fX29JMntdjf5uI3Hvr7EiZuuOfNcqrZwHm3hHCTOo6nawjlInEdTtYVzkDiPpmruOZxeb4y58ELTCv3v//6vkWTeffddr/unT59ubrzxxnM+Zs6cOUYSGxsbGxsbWxvYdu/efcFWaJWvwFyK2bNnKzc313O7sbFRBw8eVKdOneRwOHxyTLfbrfj4eO3evVtOp9Mnx/C1tnAOUts4j7ZwDhLn0Zq0hXOQ2sZ5tIVzkL6d8zDG6NChQ4qLi7vgulYZMJ07d1ZgYKBqa2u97q+trVVMTMw5HxMaGqrQ0FCv+yIiInw1ohen02n1/yCltnEOUts4j7ZwDhLn0Zq0hXOQ2sZ5tIVzkHx/Hi6X66JrWuVl1CEhIUpJSVFRUZHnvsbGRhUVFSktLc2PkwEAgNagVb4CI0m5ubm677771K9fP91444166qmndPjwYc9VSQAA4MrVagPmrrvu0v79+5WXl6eamhpdf/31KigoUHR0tL9H8wgNDdWcOXPOeuvKJm3hHKS2cR5t4RwkzqM1aQvnILWN82gL5yC1rvNwGHOx65QAAABal1b5GRgAAIALIWAAAIB1CBgAAGAdAgYAAFiHgAGAc+D6BqB1a7WXUcM3/vGPf+i5555TSUmJampqJEkxMTG66aabdP/996tLly5+nhBoHUJDQ/XBBx8oOTnZ36MAOAcuo24BtbW1+t3vfqe8vDx/j3JB7733njIyMtSuXTulp6d7fqZObW2tioqK9PXXX+tvf/ub+vXr5+dJL8/u3bs1Z84cPffcc/4e5YKOHDmi0tJSRUZGqlevXl77jh49qjVr1mjSpEl+mq7pKioqtGXLFqWlpSkpKUmffvqpli1bpmPHjumee+7RLbfc4u8RL+jM36F2pmXLlumee+5Rp06dJElPPvnktznWZTt8+LDWrFmjHTt2KDY2VhMmTPCcS2v2/vvvq2PHjurWrZsk6cUXX9TKlStVXV2thIQE5eTkaPz48X6e8uIefPBB3XnnnRo0aJC/R7ksy5cv17Zt2zR69GiNHz9eL774ohYuXKjGxkaNHTtW8+fPV1CQn14LaYFfHn3FKysrMwEBAf4e46JSU1NNVlaWaWxsPGtfY2OjycrKMgMGDPDDZC3Lhq9HZWWlSUhIMA6HwwQEBJjBgwebPXv2ePbX1NS0+nMwxpi33nrLhISEmMjISBMWFmbeeust06VLF5Oenm5uueUWExgYaIqKivw95gU5HA5z/fXXm6FDh3ptDofD9O/f3wwdOtQMGzbM32NeVHJysjlw4IAxxpjq6mqTmJhoXC6X6d+/v4mMjDRRUVHmiy++8POUF9e7d29TWFhojDHm2WefNeHh4eahhx4yK1asMA8//LBp3769+cMf/uDnKS/u9N/t7373u2bRokVm7969/h6p2RYsWGA6dOhgxo0bZ2JiYsyiRYtMp06dzM9//nPz+OOPmy5dupi8vDy/zUfANMEHH3xwwW316tVWfLMJCwszFRUV591fUVFhwsLCvsWJLs3rr79+wW3p0qWt/utx++23m8zMTLN//35TVVVlMjMzTbdu3cyXX35pjLEnYNLS0szPfvYzY4wxL7/8sunYsaP56U9/6tk/a9Ys8/3vf99f4zXJwoULTbdu3c4KraCgIFNeXu6nqZrP4XCY2tpaY4wxEydONDfddJOpq6szxhhz6NAhk56ebiZMmODPEZskPDzc7Nq1yxhjTN++fc0zzzzjtX/VqlWmV69e/hitWRwOh3n77bfNtGnTTOfOnU1wcLC59dZbzZtvvmlOnTrl7/Ga5OqrrzZ//etfjTHf/MMwMDDQ/OlPf/LsX7t2renRo4e/xiNgmuJ0STscjrO20/fb8M0mMTHRvPDCC+fd/8ILL5iEhIRvb6BLdKGvx5lfl9YsKirKfPjhh57bjY2N5oEHHjBdu3Y1n3/+uTUB43Q6TVVVlTHGmFOnTpmgoCDz/vvve/Z/9NFHJjo62l/jNdm2bdvMNddcYx555BFz/PhxY4zdAdO9e3ezYcMGr/2bN2828fHx/hitWTp16mS2b99ujPnm70lZWZnX/h07dpjw8HB/jNYsZ349jh8/blavXm0yMjJMYGCgiYuLMz/96U89f3daq/DwcM8/qowxJjg42Hz88cee27t27TLt2rXzx2jGGGO4CqkJIiMj9eyzz2rnzp1nbV988YXWrVvn7xGb5NFHH1VWVpamTZumN954Q1u3btXWrVv1xhtvaNq0aXrggQc0Y8YMf495UbGxsVq7dq0aGxvPub3//vv+HvGijhw54vW+scPh0IoVKzRmzBgNGTJEn332mR+nax6HwyFJCggIUFhYmFwul2dfhw4dVF9f76/Rmqx///4qLS3V/v371a9fP3388cee87LJ6ZmPHj2q2NhYr33/9m//pv379/tjrGYZNWqUVqxYIUkaMmSI/vKXv3jtX7NmjXr06OGP0S5ZcHCw7rzzThUUFOiLL77QlClTtGrVKvXs2dPfo11QTEyMPvnkE0lSVVWVTp065bktSeXl5YqKivLXeFyF1BQpKSnas2ePEhISzrm/rq7Oiksus7Oz1blzZy1dulS//e1vderUKUlSYGCgUlJSlJ+frzvvvNPPU15cSkqKSktLddttt51zv8PhaPVfj6SkJG3fvv2sK1yWL18uSbr11lv9MVazJSYmqqqqSldffbUkqaSkRF27dvXsr66uPusbaWvVvn17vfDCC3rllVeUnp7u+fthk+HDhysoKEhut1uVlZX63ve+59n35ZdfWvEh3sWLF+vmm2/WkCFD1K9fP/3qV7/Sxo0blZycrMrKSm3ZskWvvvqqv8e8ZF27dtXcuXM1Z84cvf322/4e54ImTpyoSZMm6bbbblNRUZFmzJihRx99VAcOHJDD4dAvfvEL3XHHHX6bj4BpggceeECHDx8+7/6uXbvq+eef/xYnunR33XWX7rrrLp04cUL/+Mc/JEmdO3dWcHCwnydruunTp1/w69GjRw/913/917c4UfP98Ic/1Msvv6x77733rH3Lly9XY2OjVq5c6YfJmmfq1Kle3+jP/IYpSW+99VarvwrpX40fP14DBw5UaWnpef/R0hrNmTPH63b79u29br/55ptWXBETFxen//mf/9GiRYv05ptvyhijbdu2affu3br55pu1efNmK66UTEhIUGBg4Hn3OxwOff/73/8WJ2q+efPmKTw8XCUlJZoyZYpmzZqlPn36aMaMGfr66681ZswYLViwwG/zcRn1JTr9l6g1/EpxAACuNATMJXI6nSorK1P37t39PQoAAFccPsR7ieg+AAD8h4ABAADWIWAu0e9+9zvPj+IHAADfLj4DAwAArMMrMAAAwDoEDAAAsA4BAwAArEPAAAAA6xAwAFqVxMREPfXUU01a63A49Nprr13W8fLz8xUREXFZzwHg28fvQgJgrb1796pjx47+HgOAHxAwAFrc8ePHFRIS4vPjxMTE+PwYAFon3kICcFFDhw5VTk6OcnJy5HK51LlzZ/3nf/6n51dqJCYmasGCBZo0aZKcTqeysrIkSf/93/+tQYMGKTw8XPHx8XrooYe8fpP4vn37NGbMGIWHh6tbt25atWpVs+Y68y2kXbt2yeFwaO3atRo2bJjatWunPn36qKSkxOsx+fn56tq1q9q1a6cf/vCHOnDgwFnP+/rrr+uGG25QWFiYunfvrnnz5unkyZOSpPnz5ysuLs7rcZmZmRo2bJgaGxubNT+Ay2AA4CKGDBli2rdvb6ZNm2Y+/fRT86c//cm0a9fOPPPMM8YYYxISEozT6TRPPPGE2bFjh2e76qqrzNKlS81nn31mNm/ebPr27Wvuv/9+z/OOGjXK9OnTx5SUlJjt27ebm266yYSHh5ulS5c2aS5J5tVXXzXGGLNz504jySQlJZl169aZyspKc8cdd5iEhARz4sQJY4wxW7ZsMQEBAWbx4sWmsrLSLFu2zERERBiXy+V5zuLiYuN0Ok1+fr75/PPPzYYNG0xiYqKZO3euMcaYkydPmrS0NHP77bcbY4xZvny5iYiIMF9++eVl/lcG0BwEDICLGjJkiElOTjaNjY2e+2bOnGmSk5ONMd8EzOlv6KdNnjzZZGVled3397//3QQEBJgjR46YyspKI8ls27bNs7+iosJIuqyA+f3vf+/ZX15ebiSZiooKY4wxEyZMMKNHj/Z6jrvuussrYIYPH24ef/xxrzUvvviiiY2N9dz+/PPPTYcOHczMmTNNeHi4WbVqVZPmBdByeAsJQJMMGDBADofDczstLU1VVVU6deqUJKlfv35e6z/44APl5+erffv2ni0jI0ONjY3auXOnKioqFBQUpJSUFM9jkpKSLvuKoN69e3v+HBsbK+mbt6okqaKiQqmpqV7r09LSzpp7/vz5XnNPmTJFe/fu1ddffy1J6t69u5544gktXrxYt956q+6+++7LmhlA8/EhXgAt4qqrrvK63dDQoB//+Md66KGHzlrbtWtXffbZZz6ZIzg42PPn08HVnM+mNDQ0aN68eRo7duxZ+8LCwjx/Li4uVmBgoHbt2qWTJ08qKIj/OwW+TfyNA9AkW7du9bq9ZcsWffe731VgYOA5199www365JNP1KNHj3PuT0pK0smTJ1VaWqr+/ftLkiorK1VXV9eic58pOTn5nOdxphtuuEGVlZXnnVuSVq9erbVr12rjxo268847tWDBAs2bN88nMwM4N95CAtAk1dXVys3NVWVlpV5++WX95je/0bRp0867fubMmXr33XeVk5OjsrIyVVVV6fXXX1dOTo4kqWfPnho5cqR+/OMfa+vWrSotLdWPfvQjhYeH++wcHnroIRUUFOiJJ55QVVWVli9froKCAq81eXl5+uMf/6h58+apvLxcFRUVeuWVV/TYY49Jkr766itNnTpVixcv1sCBA/X888/r8ccfPyuEAPgWAQOgSSZNmqQjR47oxhtvVHZ2tqZNm+a5XPpcevfurU2bNumzzz7ToEGD1LdvX+Xl5SkuLs6z5vnnn1dcXJyGDBmisWPHKisrS1FRUT47hwEDBujZZ5/VsmXL1KdPH23YsMETJqdlZGRo3bp12rBhg/r3768BAwZo6dKlSkhIkDFG999/v2688UZPiGVkZGjq1Km655571NDQ4LPZAXhzGPP/f5ADAJzH0KFDdf311zf5R/wDgK/xCgwAALAOAQOgVVq1apXXpcxnbtdee62/xwPgZ7yFBKBVOnTokGpra8+5Lzg4WAkJCd/yRABaEwIGAABYh7eQAACAdQgYAABgHQIGAABYh4ABAADWIWAAAIB1CBgAAGAdAgYAAFjn/wF5rnwSZRRNNQAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ea_df.pred_index.value_counts().sort_index().iloc[0:10].plot.bar()"
   ]
  }
 ],
 "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
}