Experiments/kaggle/NBME/src/fast-awp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a8caa43a",
   "metadata": {
    "papermill": {
     "duration": 0.014641,
     "end_time": "2022-04-06T12:09:18.139518",
     "exception": false,
     "start_time": "2022-04-06T12:09:18.124877",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Fast AWP with small overhead\n",
    "\n",
    "The *Adversarial Weight Perturbation (AWP)* was used in the [first-place solution of the Feedback Prize](https://www.kaggle.com/competitions/feedback-prize-2021/discussion/313177) (@wht1996) and it is showen to be  [effective in this competition as well](https://www.kaggle.com/competitions/nbme-score-clinical-patient-notes/discussion/315707) (@hengck23). The perturbation, however, requires two gradient computations per optimization step, and therefore takes twice more computation time (correct me if I am wrong). This implementation removes the overhead by approximating the gradient with the running mean stored in the Adam optimizer. \n",
    "\n",
    "Since the gradient given to the optimizer is perturbed by\n",
    "AWP, the AWP in this notebook uses \"wrong\" gradients, and may not work as intended. Still, the factor of two in time is large and seems to have some good effects; maybe AWP works with inaccurate gradients. Let me know in the comment if you see this AWP underperforms compared to the original. My PyTorch coding is pretty random and comments are welcome on that aspect, too, e.g., detach and copy are not appropriate. Maybe my code doing something completely random.\n",
    "\n",
    "The removal of additional gradient calculation also allows gradient accumulation in the usual way. The overhead in computation time is little, but overhead in GPU RAM exists for a copy of model weights (0.5 - 1 GB).\n",
    "\n",
    "\n",
    "## Parameters\n",
    "\n",
    "Two parameters control the amount of perturbations in units of fractional changes in the weights; 0.01 means 1% perturbation in weights (model parameters).\n",
    "\n",
    "```\n",
    "adv_lr or γ: fractional change in weight along the direction of gradient\n",
    "adv_eps or ε: the change in weight is limited to this fraction\n",
    "```\n",
    "\n",
    "Two parameters are similar to learning rate and max norm in gradient clipping, but in units of fraction of weights.\n",
    "\n",
    "The AWP gradients are evaluated at perturbed location:\n",
    "\n",
    "$$ w \\mapsto w + \\delta w = w + \\gamma \\frac{\\nabla}{\\lVert{\\nabla} \\rVert} \\lVert w \\rVert, $$\n",
    "\n",
    "but change in each component is limitted to,\n",
    "\n",
    "$$ |\\delta w_i| \\le \\epsilon |w_i|. $$\n",
    "\n",
    "where `∇ = param.grad` is the gradient for the weight `w = param.data`.\n",
    "\n",
    "\n",
    "You only need to see `class AWP` and use it in your training loop; you do not need to read other parts of my code.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## Reference\n",
    "\n",
    "This AWP modifies the 1st-place code in the Feedback Prize solution by @wht1996:\n",
    "\n",
    "https://www.kaggle.com/code/wht1996/feedback-nn-train \n",
    "\n",
    "\n",
    "The original paper is,\n",
    "\n",
    "Wu, Xia, and  Wang (2020), Adversarial Weight Perturbation Helps Robust Generalization\n",
    "https://arxiv.org/abs/2004.05884\n",
    "\n",
    "Authors' implementation: https://github.com/csdongxian/AWP\n",
    "\n",
    "The model and many codes for training are borrowed from the Nakama baseline >ω</ Thanks!:\n",
    "\n",
    "https://www.kaggle.com/code/yasufuminakama/nbme-deberta-base-baseline-train/notebook\n",
    "\n",
    "And many thanks to hengck23 for sharing the usefulness of AWP and the working parameters:\n",
    "\n",
    "https://www.kaggle.com/competitions/nbme-score-clinical-patient-notes/discussion/315707"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5ac9b2a1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:18.181055Z",
     "iopub.status.busy": "2022-04-06T12:09:18.180275Z",
     "iopub.status.idle": "2022-04-06T12:09:24.491789Z",
     "shell.execute_reply": "2022-04-06T12:09:24.492225Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.4963Z"
    },
    "papermill": {
     "duration": 6.33775,
     "end_time": "2022-04-06T12:09:24.492524",
     "exception": false,
     "start_time": "2022-04-06T12:09:18.154774",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import os\n",
    "import ast\n",
    "import time\n",
    "import yaml\n",
    "import random\n",
    "import argparse\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import AdamW\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.model_selection import StratifiedGroupKFold\n",
    "\n",
    "import transformers\n",
    "from transformers import AutoTokenizer, AutoConfig, AutoModel\n",
    "from transformers import get_cosine_schedule_with_warmup\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "transformers.logging.set_verbosity_error()\n",
    "\n",
    "def seed_everything(seed=42):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    \n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "811bc3bf",
   "metadata": {
    "papermill": {
     "duration": 0.012162,
     "end_time": "2022-04-06T12:09:24.517262",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.505100",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Class AWP\n",
    "\n",
    "Modified from https://www.kaggle.com/code/wht1996/feedback-nn-train (wht1996)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9a62ff6a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:24.545208Z",
     "iopub.status.busy": "2022-04-06T12:09:24.544384Z",
     "iopub.status.idle": "2022-04-06T12:09:24.557413Z",
     "shell.execute_reply": "2022-04-06T12:09:24.557029Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.521294Z"
    },
    "papermill": {
     "duration": 0.028008,
     "end_time": "2022-04-06T12:09:24.557555",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.529547",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class AWP:\n",
    "    def __init__(self, model, optimizer, *, adv_param='weight',\n",
    "                 adv_lr=0.001, adv_eps=0.001):\n",
    "        self.model = model\n",
    "        self.optimizer = optimizer\n",
    "        self.adv_param = adv_param\n",
    "        self.adv_lr = adv_lr\n",
    "        self.adv_eps = adv_eps\n",
    "        self.backup = {}\n",
    "\n",
    "    def perturb(self, input_ids, attention_mask, y, criterion):\n",
    "        \"\"\"\n",
    "        Perturb model parameters for AWP gradient\n",
    "        Call before loss and loss.backward()\n",
    "        \"\"\"\n",
    "        self._save()  # save model parameters\n",
    "        self._attack_step()  # perturb weights\n",
    "\n",
    "    def _attack_step(self):\n",
    "        e = 1e-6\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and param.grad is not None and self.adv_param in name:\n",
    "                grad = self.optimizer.state[param]['exp_avg']\n",
    "                norm_grad = torch.norm(grad)\n",
    "                norm_data = torch.norm(param.detach())\n",
    "\n",
    "                if norm_grad != 0 and not torch.isnan(norm_grad):\n",
    "                    # Set lower and upper limit in change\n",
    "                    limit_eps = self.adv_eps * param.detach().abs()\n",
    "                    param_min = param.data - limit_eps\n",
    "                    param_max = param.data + limit_eps\n",
    "\n",
    "                    # Perturb along gradient\n",
    "                    # w += (adv_lr * |w| / |grad|) * grad\n",
    "                    param.data.add_(grad, alpha=(self.adv_lr * (norm_data + e) / (norm_grad + e)))\n",
    "\n",
    "                    # Apply the limit to the change\n",
    "                    param.data.clamp_(param_min, param_max)\n",
    "\n",
    "    def _save(self):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and param.grad is not None and self.adv_param in name:\n",
    "                if name not in self.backup:\n",
    "                    self.backup[name] = param.clone().detach()\n",
    "                else:\n",
    "                    self.backup[name].copy_(param.data)\n",
    "\n",
    "    def restore(self):\n",
    "        \"\"\"\n",
    "        Restore model parameter to correct position; AWP do not perturbe weights, it perturb gradients\n",
    "        Call after loss.backward(), before optimizer.step()\n",
    "        \"\"\"\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if name in self.backup:\n",
    "                param.data.copy_(self.backup[name])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a15c0b7",
   "metadata": {
    "papermill": {
     "duration": 0.012494,
     "end_time": "2022-04-06T12:09:24.582327",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.569833",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Changes in AWP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65c82d87",
   "metadata": {
    "papermill": {
     "duration": 0.012007,
     "end_time": "2022-04-06T12:09:24.606745",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.594738",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "- grad is using Adam exponentially averaged gradient;\n",
    "  * this saves computing gradient twice and **speeds up by factor of 2**;\n",
    "  * but this grad is pertubed by AWP and not the true gradient.\n",
    "- adv_step is removed; since grads are not computed, there is no reason to do multiple steps.\n",
    "- since adv_step=1, weight eps ranges do not have to be saved,\n",
    "  * which reduce GPU RAM from 3 extra copies of wegiths to 1 copy.\n",
    "\n",
    "\n",
    "## Original code by wht1996\n",
    "\n",
    "```python\n",
    "class AWP:\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        optimizer,\n",
    "        adv_param=\"weight\",\n",
    "        adv_lr=1,\n",
    "        adv_eps=0.2,\n",
    "        start_epoch=0,\n",
    "        adv_step=1,\n",
    "        scaler=None\n",
    "    ):\n",
    "        self.model = model\n",
    "        self.optimizer = optimizer\n",
    "        self.adv_param = adv_param\n",
    "        self.adv_lr = adv_lr\n",
    "        self.adv_eps = adv_eps\n",
    "        self.start_epoch = start_epoch\n",
    "        self.adv_step = adv_step\n",
    "        self.backup = {}\n",
    "        self.backup_eps = {}\n",
    "        self.scaler = scaler\n",
    "  \n",
    "    def attack_backward(self, x, y, attention_mask,epoch):\n",
    "        if (self.adv_lr == 0) or (epoch < self.start_epoch):\n",
    "            return None\n",
    "\n",
    "        self._save() \n",
    "        for i in range(self.adv_step):\n",
    "            self._attack_step() \n",
    "            with torch.cuda.amp.autocast():\n",
    "                adv_loss, tr_logits = self.model(input_ids=x, attention_mask=attention_mask, labels=y)\n",
    "                adv_loss = adv_loss.mean()\n",
    "            self.optimizer.zero_grad()\n",
    "            self.scaler.scale(adv_loss).backward()\n",
    "            \n",
    "        self._restore()\n",
    "\n",
    "    def _attack_step(self):\n",
    "        e = 1e-6\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and param.grad is not None and self.adv_param in name:\n",
    "                norm1 = torch.norm(param.grad)\n",
    "                norm2 = torch.norm(param.data.detach())\n",
    "                if norm1 != 0 and not torch.isnan(norm1):\n",
    "                    r_at = self.adv_lr * param.grad / (norm1 + e) * (norm2 + e)\n",
    "                    param.data.add_(r_at)\n",
    "                    param.data = torch.min(\n",
    "                        torch.max(param.data, self.backup_eps[name][0]), self.backup_eps[name][1]\n",
    "                    )\n",
    "                # param.data.clamp_(*self.backup_eps[name])\n",
    "\n",
    "    def _save(self):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and param.grad is not None and self.adv_param in name:\n",
    "                if name not in self.backup:\n",
    "                    self.backup[name] = param.data.clone()\n",
    "                    grad_eps = self.adv_eps * param.abs().detach()\n",
    "                    self.backup_eps[name] = (\n",
    "                        self.backup[name] - grad_eps,\n",
    "                        self.backup[name] + grad_eps,\n",
    "                    )\n",
    "\n",
    "    def _restore(self,):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if name in self.backup:\n",
    "                param.data = self.backup[name]\n",
    "        self.backup = {}\n",
    "        self.backup_eps = {}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b03ad765",
   "metadata": {
    "papermill": {
     "duration": 0.012105,
     "end_time": "2022-04-06T12:09:24.631736",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.619631",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f8aa06fa",
   "metadata": {
    "_kg_hide-input": true,
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:24.682013Z",
     "iopub.status.busy": "2022-04-06T12:09:24.681114Z",
     "iopub.status.idle": "2022-04-06T12:09:24.682852Z",
     "shell.execute_reply": "2022-04-06T12:09:24.683329Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.614584Z"
    },
    "papermill": {
     "duration": 0.03956,
     "end_time": "2022-04-06T12:09:24.683450",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.643890",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def nakama_fix_annotations(train):\n",
    "    # Fix incorrect annotations (from Nakama baseline)\n",
    "    train.loc[338, 'annotation'] = \"['father heart attack']\"\n",
    "    train.loc[338, 'location'] = \"['764 783']\"\n",
    "\n",
    "    train.loc[621, 'annotation'] = \"['for the last 2-3 months']\"\n",
    "    train.loc[621, 'location'] = \"['77 100']\"\n",
    "\n",
    "    train.loc[655, 'annotation'] = \"['no heat intolerance'], ['no cold intolerance']\"\n",
    "    train.loc[655, 'location'] = \"['285 292;301 312', '285 287;296 312']\"\n",
    "\n",
    "    train.loc[1262, 'annotation'] = \"['mother thyroid problem']\"\n",
    "    train.loc[1262, 'location'] = \"['551 557;565 580']\"\n",
    "\n",
    "    train.loc[1265, 'annotation'] = \"[\\\"felt like he was going to 'pass out'\\\"]\"\n",
    "    train.loc[1265, 'location'] = \"['131 135;181 212']\"\n",
    "\n",
    "    train.loc[1396, 'annotation'] = \"['stool , with no blood']\"\n",
    "    train.loc[1396, 'location'] = \"['259 280']\"\n",
    "\n",
    "    train.loc[1591, 'annotation'] = \"['diarrhoe non blooody']\"\n",
    "    train.loc[1591, 'location'] = \"['176 184;201 212']\"\n",
    "\n",
    "    train.loc[1615, 'annotation'] = \"['diarrhea for last 2-3 days']\"\n",
    "    train.loc[1615, 'location'] = \"['249 257;271 288']\"\n",
    "\n",
    "    train.loc[1664, 'annotation'] = \"['no vaginal discharge']\"\n",
    "    train.loc[1664, 'location'] = \"['822 824;907 924']\"\n",
    "\n",
    "    train.loc[1714, 'annotation'] = \"['started about 8-10 hours ago']\"\n",
    "    train.loc[1714, 'location'] = \"['101 129']\"\n",
    "\n",
    "    train.loc[1929, 'annotation'] = \"['no blood in the stool']\"\n",
    "    train.loc[1929, 'location'] = \"['531 539;549 561']\"\n",
    "\n",
    "    train.loc[2134, 'annotation'] = \"['last sexually active 9 months ago']\"\n",
    "    train.loc[2134, 'location'] = \"['540 560;581 593']\"\n",
    "\n",
    "    train.loc[2191, 'annotation'] = \"['right lower quadrant pain']\"\n",
    "    train.loc[2191, 'location'] = \"['32 57']\"\n",
    "\n",
    "    train.loc[2553, 'annotation'] = \"['diarrhoea no blood']\"\n",
    "    train.loc[2553, 'location'] = \"['308 317;376 384']\"\n",
    "\n",
    "    train.loc[3124, 'annotation'] = \"['sweating']\"\n",
    "    train.loc[3124, 'location'] = \"['549 557']\"\n",
    "\n",
    "    train.loc[3858, 'annotation'] = \"['previously as regular', 'previously eveyr 28-29 days', 'previously lasting 5 days'], 'previously regular flow']\"\n",
    "    train.loc[3858, 'location'] = \"['102 123', '102 112;125 141', '102 112;143 157', '102 112;159 171']\"\n",
    "\n",
    "    train.loc[4373, 'annotation'] = \"['for 2 months']\"\n",
    "    train.loc[4373, 'location'] = \"['33 45']\"\n",
    "\n",
    "    train.loc[4763, 'annotation'] = \"['35 year old']\"\n",
    "    train.loc[4763, 'location'] = \"['5 16']\"\n",
    "\n",
    "    train.loc[4782, 'annotation'] = \"['darker brown stools']\"\n",
    "    train.loc[4782, 'location'] = \"['175 194']\"\n",
    "\n",
    "    train.loc[4908, 'annotation'] = \"['uncle with peptic ulcer']\"\n",
    "    train.loc[4908, 'location'] = \"['700 723']\"\n",
    "\n",
    "    train.loc[6016, 'annotation'] = \"['difficulty falling asleep']\"\n",
    "    train.loc[6016, 'location'] = \"['225 250']\"\n",
    "\n",
    "    train.loc[6192, 'annotation'] = \"['helps to take care of aging mother and in-laws']\"\n",
    "    train.loc[6192, 'location'] = \"['197 218;236 260']\"\n",
    "\n",
    "    train.loc[6380, 'annotation'] = \"['No hair changes', 'No skin changes', 'No GI changes', 'No palpitations', 'No excessive sweating']\"\n",
    "    train.loc[6380, 'location'] = \"['480 482;507 519', '480 482;499 503;512 519', '480 482;521 531', '480 482;533 545', '480 482;564 582']\"\n",
    "\n",
    "    train.loc[6562, 'annotation'] = \"['stressed due to taking care of her mother', 'stressed due to taking care of husbands parents']\"\n",
    "    train.loc[6562, 'location'] = \"['290 320;327 337', '290 320;342 358']\"\n",
    "\n",
    "    train.loc[6862, 'annotation'] = \"['stressor taking care of many sick family members']\"\n",
    "    train.loc[6862, 'location'] = \"['288 296;324 363']\"\n",
    "\n",
    "    train.loc[7022, 'annotation'] = \"['heart started racing and felt numbness for the 1st time in her finger tips']\"\n",
    "    train.loc[7022, 'location'] = \"['108 182']\"\n",
    "\n",
    "    train.loc[7422, 'annotation'] = \"['first started 5 yrs']\"\n",
    "    train.loc[7422, 'location'] = \"['102 121']\"\n",
    "\n",
    "    train.loc[8876, 'annotation'] = \"['No shortness of breath']\"\n",
    "    train.loc[8876, 'location'] = \"['481 483;533 552']\"\n",
    "\n",
    "    train.loc[9027, 'annotation'] = \"['recent URI', 'nasal stuffines, rhinorrhea, for 3-4 days']\"\n",
    "    train.loc[9027, 'location'] = \"['92 102', '123 164']\"\n",
    "\n",
    "    train.loc[9938, 'annotation'] = \"['irregularity with her cycles', 'heavier bleeding', 'changes her pad every couple hours']\"\n",
    "    train.loc[9938, 'location'] = \"['89 117', '122 138', '368 402']\"\n",
    "\n",
    "    train.loc[9973, 'annotation'] = \"['gaining 10-15 lbs']\"\n",
    "    train.loc[9973, 'location'] = \"['344 361']\"\n",
    "\n",
    "    train.loc[10513, 'annotation'] = \"['weight gain', 'gain of 10-16lbs']\"\n",
    "    train.loc[10513, 'location'] = \"['600 611', '607 623']\"\n",
    "\n",
    "    train.loc[11551, 'annotation'] = \"['seeing her son knows are not real']\"\n",
    "    train.loc[11551, 'location'] = \"['386 400;443 461']\"\n",
    "\n",
    "    train.loc[11677, 'annotation'] = \"['saw him once in the kitchen after he died']\"\n",
    "    train.loc[11677, 'location'] = \"['160 201']\"\n",
    "\n",
    "    train.loc[12124, 'annotation'] = \"['tried Ambien but it didnt work']\"\n",
    "    train.loc[12124, 'location'] = \"['325 337;349 366']\"\n",
    "\n",
    "    train.loc[12279, 'annotation'] = \"['heard what she described as a party later than evening these things did not actually happen']\"\n",
    "    train.loc[12279, 'location'] = \"['405 459;488 524']\"\n",
    "\n",
    "    train.loc[12289, 'annotation'] = \"['experienced seeing her son at the kitchen table these things did not actually happen']\"\n",
    "    train.loc[12289, 'location'] = \"['353 400;488 524']\"\n",
    "\n",
    "    train.loc[13238, 'annotation'] = \"['SCRACHY THROAT', 'RUNNY NOSE']\"\n",
    "    train.loc[13238, 'location'] = \"['293 307', '321 331']\"\n",
    "\n",
    "    train.loc[13297, 'annotation'] = \"['without improvement when taking tylenol', 'without improvement when taking ibuprofen']\"\n",
    "    train.loc[13297, 'location'] = \"['182 221', '182 213;225 234']\"\n",
    "\n",
    "    train.loc[13299, 'annotation'] = \"['yesterday', 'yesterday']\"\n",
    "    train.loc[13299, 'location'] = \"['79 88', '409 418']\"\n",
    "\n",
    "    train.loc[13845, 'annotation'] = \"['headache global', 'headache throughout her head']\"\n",
    "    train.loc[13845, 'location'] = \"['86 94;230 236', '86 94;237 256']\"\n",
    "\n",
    "    train.loc[14083, 'annotation'] = \"['headache generalized in her head']\"\n",
    "    train.loc[14083, 'location'] = \"['56 64;156 179']\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2a625db3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:24.731576Z",
     "iopub.status.busy": "2022-04-06T12:09:24.721207Z",
     "iopub.status.idle": "2022-04-06T12:09:24.733892Z",
     "shell.execute_reply": "2022-04-06T12:09:24.733454Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.645515Z"
    },
    "papermill": {
     "duration": 0.037825,
     "end_time": "2022-04-06T12:09:24.733999",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.696174",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def _replace_feature_text(text):\n",
    "    text = text.replace('I-year', '1-year')\n",
    "    text = text.replace('-OR-', ' or ')\n",
    "    text = text.replace('-', ' ')\n",
    "\n",
    "    return text\n",
    "\n",
    "def get_segments(locations):\n",
    "    \"\"\"\n",
    "    Parse location list to sorted list of segments\n",
    "\n",
    "    Args:\n",
    "      location (str): ['85 99', '126 138', '126 131;143 151']\n",
    "\n",
    "    Returns: list[tuple]\n",
    "      List of (begin, end)\n",
    "    \"\"\"\n",
    "    segs = []\n",
    "\n",
    "    assert isinstance(locations, str)\n",
    "\n",
    "    locations = ast.literal_eval(locations)  # str -> list[str]\n",
    "    for loc in locations:\n",
    "        for span in loc.split(';'):\n",
    "            segs.append(tuple(map(int, span.split(' '))))\n",
    "\n",
    "    segs.sort(key=lambda pair: pair[0])\n",
    "\n",
    "    return segs\n",
    "\n",
    "\n",
    "def create_label(segments, offset_mapping):\n",
    "    n = len(offset_mapping)\n",
    "    y = np.zeros(n, dtype=np.float32)\n",
    "\n",
    "    if not segments:\n",
    "        return y\n",
    "\n",
    "    iseg = 0\n",
    "    seg = segments[iseg]\n",
    "    nseg = len(segments)\n",
    "\n",
    "    k = 0\n",
    "    while k < n:\n",
    "        begin, end = offset_mapping[k]\n",
    "        if end <= seg[0]:  # token is left of seg\n",
    "            k += 1\n",
    "            continue\n",
    "        elif begin < seg[1]:  # token overlaps with seg\n",
    "            y[k] = 1\n",
    "            k += 1\n",
    "        else:  # begin passed seg: seg[1] <= begin\n",
    "            iseg += 1\n",
    "            if iseg < nseg:\n",
    "                seg = segments[iseg]\n",
    "            else:\n",
    "                break  # All segments processed\n",
    "\n",
    "    return y\n",
    "\n",
    "\n",
    "def exclude_feature_label(input_ids, y, sep):\n",
    "    \"\"\"\n",
    "    Set y = -1 for feature_text; compute loss only with note texts excluding feature texts\n",
    "\n",
    "    Result:\n",
    "      y modified to -1 after first [SEP]\n",
    "    \"\"\"\n",
    "    n = len(input_ids)\n",
    "    assert n >= 2 and input_ids[-1] == sep\n",
    "    y[-1] = -1\n",
    "\n",
    "    for k in range(n - 2, -1, -1):\n",
    "        if input_ids[k] == sep:\n",
    "            break\n",
    "        else:\n",
    "            y[k] = -1\n",
    "\n",
    "\n",
    "def create_character_indices(segments):\n",
    "    \"\"\"\n",
    "    Character indices\n",
    "\n",
    "    Args:\n",
    "      segments (list[tuple]): List of (begin, end)\n",
    "    \"\"\"\n",
    "    s = set()\n",
    "\n",
    "    for begin, end in segments:\n",
    "        s.update(range(begin, end))\n",
    "\n",
    "    return sorted(list(s))\n",
    "\n",
    "\n",
    "def create_data(train, tokenizer, *, max_length=1024, pbar=False):\n",
    "    \"\"\"\n",
    "    Create input_ids and label array y\n",
    "\n",
    "    Args:\n",
    "      train (pd.DataFrame)\n",
    "      tokenizer (str or tokenizer): path to tokenizer dir if str\n",
    "\n",
    "    Returns: list[dict]\n",
    "      input_ids (np.array[int])\n",
    "      n (int): number of tokens\n",
    "      y (np.array[float32]): binary annotation\n",
    "    \"\"\"\n",
    "    sep = tokenizer.sep_token_id\n",
    "    nsep = 2\n",
    "\n",
    "    annotated = 'location' in train.columns\n",
    "\n",
    "    data = []\n",
    "    for i, r in tqdm(train.iterrows(), disable=(not pbar), total=len(train)):\n",
    "        text = r.pn_history\n",
    "        feature_text = r.feature_text\n",
    "\n",
    "        o = tokenizer(text, feature_text,\n",
    "                      add_special_tokens=True, max_length=max_length,\n",
    "                      truncation=True,\n",
    "                      return_offsets_mapping=True)\n",
    "\n",
    "        # Input ids\n",
    "        input_ids = o['input_ids']\n",
    "        n = len(input_ids)\n",
    "\n",
    "        input_ids = np.array(o['input_ids'], dtype=np.int32)\n",
    "        assert np.sum(input_ids == sep) == nsep  # Two tokens seperated by [SEP], <s/><s/> for roberta\n",
    "\n",
    "        # Attention mask\n",
    "        attention_mask = np.array(o['attention_mask'])\n",
    "        assert np.all(attention_mask == 1)\n",
    "\n",
    "        # Label\n",
    "        if annotated:\n",
    "            segs = get_segments(r['location'])\n",
    "            y = create_label(segs, o['offset_mapping'])\n",
    "\n",
    "            exclude_feature_label(input_ids, y, sep)\n",
    "\n",
    "            # Character label\n",
    "            label = create_character_indices(segs)\n",
    "        else:\n",
    "            y = None\n",
    "            label = None\n",
    "\n",
    "        d = {'id': r['id'],\n",
    "             'input_ids': input_ids,\n",
    "             'text': text,\n",
    "             'label': label,\n",
    "             'y': y,\n",
    "             'n': n,\n",
    "             'offset_mapping': o['offset_mapping']}\n",
    "        data.append(d)\n",
    "\n",
    "    return np.array(data)\n",
    "\n",
    "\n",
    "class Dataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    Dataset(data)\n",
    "      data (np.array or list-like): input_ids and y\n",
    "    \"\"\"\n",
    "    def __init__(self, data, *, max_length=512):\n",
    "        self.data = data\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        d = self.data[i]\n",
    "        n = min(d['n'], self.max_length)\n",
    "\n",
    "        input_ids = np.zeros(self.max_length, dtype=int)\n",
    "        input_ids[:n] = d['input_ids']\n",
    "\n",
    "        attention_mask = np.zeros(self.max_length, dtype=int)\n",
    "        attention_mask[:n] = 1\n",
    "\n",
    "        y = np.full(self.max_length, -1, dtype=np.float32)\n",
    "        y[:n] = d['y']\n",
    "\n",
    "        return {'input_ids': input_ids,\n",
    "                'attention_mask': attention_mask,\n",
    "                'y': y, 'n': n}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "748b495e",
   "metadata": {
    "papermill": {
     "duration": 0.012064,
     "end_time": "2022-04-06T12:09:24.758201",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.746137",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fe84252f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:24.785297Z",
     "iopub.status.busy": "2022-04-06T12:09:24.784473Z",
     "iopub.status.idle": "2022-04-06T12:09:24.792751Z",
     "shell.execute_reply": "2022-04-06T12:09:24.792315Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.675004Z"
    },
    "papermill": {
     "duration": 0.022466,
     "end_time": "2022-04-06T12:09:24.792860",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.770394",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, model_dir, *, dropout=0.2, pretrained=True):\n",
    "        super().__init__()\n",
    "\n",
    "        # Transformer\n",
    "        config = AutoConfig.from_pretrained(model_dir, add_pooling_layer=False)\n",
    "        if pretrained:\n",
    "            self.transformer = AutoModel.from_pretrained(model_dir, config=config)\n",
    "        else:\n",
    "            self.transformer = AutoModel.from_config(config)\n",
    "\n",
    "        self.fc_dropout = nn.Dropout(dropout)\n",
    "        self.fc = nn.Linear(config.hidden_size, 1)\n",
    "\n",
    "        self._init_weights(self.fc, config)\n",
    "\n",
    "    def _init_weights(self, module, config):\n",
    "        module.weight.data.normal_(mean=0.0, std=config.initializer_range)\n",
    "        if module.bias is not None:\n",
    "            module.bias.data.zero_()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        out = self.transformer(input_ids, attention_mask)\n",
    "        x = out['last_hidden_state']  # batch_size x max_length (512) x 768\n",
    "\n",
    "        x = self.fc_dropout(x)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e31b132",
   "metadata": {
    "papermill": {
     "duration": 0.012036,
     "end_time": "2022-04-06T12:09:24.817152",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.805116",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Train and evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6a73f988",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:24.861867Z",
     "iopub.status.busy": "2022-04-06T12:09:24.860978Z",
     "iopub.status.idle": "2022-04-06T12:09:24.862850Z",
     "shell.execute_reply": "2022-04-06T12:09:24.863231Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.6892Z"
    },
    "papermill": {
     "duration": 0.034043,
     "end_time": "2022-04-06T12:09:24.863352",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.829309",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_prediction(y_pred, d, *, th=0.5):\n",
    "    \"\"\"\n",
    "    Create character-level prediction\n",
    "\n",
    "    Args:\n",
    "      pred: pred['y_pred'] token-level prediction\n",
    "      d: d['offset_mapping']\n",
    "    \"\"\"\n",
    "    text = d['text']\n",
    "    offset_mapping = d['offset_mapping']\n",
    "\n",
    "    # Map token-level prob to character-level prob\n",
    "    y_prob = np.zeros(len(text))  # character-wise probabilities\n",
    "    i = 0\n",
    "    for p, (begin, end) in zip(y_pred, offset_mapping):\n",
    "        if i > 0 and begin == 0 and end == 0:\n",
    "            break  # This is end of patient note [sep]\n",
    "\n",
    "        y_prob[i:end] = p    # Set space before begin with p too (deberta style)\n",
    "        i = end\n",
    "\n",
    "    i_begin = i_last = None\n",
    "    li = []\n",
    "    for i, (x, p) in enumerate(zip(text, y_prob)):\n",
    "        if p >= th:\n",
    "            if i_begin is None and x != ' ':  # Do not include first space in span\n",
    "                i_begin = i_last = i\n",
    "                li.append(i)\n",
    "            elif i_begin is not None:         # Positive token is continuing\n",
    "                assert i_last + 1 == i\n",
    "                i_last = i\n",
    "                li.append(i)\n",
    "        else:\n",
    "            i_begin = i_last = None           # Negative\n",
    "\n",
    "    d = {'text': text,\n",
    "         'y_prob': y_prob,\n",
    "         'indices': li}\n",
    "\n",
    "    return d\n",
    "\n",
    "\n",
    "def compute_score(preds, data):\n",
    "    \"\"\"\n",
    "    TP, FN, FP are collected globally and\n",
    "    f1 score is computed at the end\n",
    "    \"\"\"\n",
    "    assert len(preds) == len(data)\n",
    "\n",
    "    tp = 0\n",
    "    denom = 0\n",
    "    for pred, d in zip(preds, data):\n",
    "        c = create_prediction(pred['y_pred'], d)\n",
    "        x = set(d['label'])\n",
    "        y = set(c['indices'])\n",
    "        tp += len(x.intersection(y))\n",
    "        denom += len(x) + len(y)\n",
    "\n",
    "    return 2 * tp / denom\n",
    "\n",
    "\n",
    "def evaluate(model, loader, criterion):\n",
    "    tb = time.time()\n",
    "    was_training = model.training\n",
    "    model.eval()\n",
    "\n",
    "    n_sum = 0\n",
    "    loss_sum = 0.0\n",
    "    preds = []\n",
    "    for d in loader:\n",
    "        input_ids = d['input_ids'].to(device)\n",
    "        attention_mask = d['attention_mask'].to(device)\n",
    "        y = d['y'].to(device)\n",
    "        n = y.size(0)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            y_pred = model(input_ids, attention_mask)\n",
    "\n",
    "        loss = criterion(y_pred.view(-1), y.view(-1))\n",
    "        loss = torch.masked_select(loss, y.view(-1) != -1).mean()\n",
    "\n",
    "        n_sum += n\n",
    "        loss_sum += n * loss.item()\n",
    "\n",
    "        y_pred = y_pred.sigmoid().cpu().numpy()\n",
    "        y = d['y'].numpy()\n",
    "\n",
    "        for k, m in enumerate(d['n']):\n",
    "            preds.append({'y_pred': y_pred[k, :m].copy(),\n",
    "                          'y': y[k, :m].copy()})\n",
    "\n",
    "        del loss, y_pred, input_ids, attention_mask, y\n",
    "\n",
    "    model.train(was_training)\n",
    "\n",
    "    val = {'time': time.time() - tb,\n",
    "           'loss': loss_sum / n_sum,\n",
    "           'preds': preds}\n",
    "\n",
    "    return val\n",
    "\n",
    "\n",
    "def get_optimizer_params(model, encoder_lr, decoder_lr, weight_decay=0.0):\n",
    "    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
    "    optimizer_parameters = [\n",
    "        {'params': [p for n, p in model.transformer.named_parameters() if not any(nd in n for nd in no_decay)],\n",
    "         'lr': encoder_lr, 'weight_decay': weight_decay},\n",
    "        {'params': [p for n, p in model.transformer.named_parameters() if any(nd in n for nd in no_decay)],\n",
    "         'lr': encoder_lr, 'weight_decay': 0.0},\n",
    "        {'params': [p for n, p in model.named_parameters() if 'transformer' not in n],\n",
    "         'lr': decoder_lr, 'weight_decay': 0.0}\n",
    "    ]\n",
    "\n",
    "    return optimizer_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "51f475f8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T12:09:24.916163Z",
     "iopub.status.busy": "2022-04-06T12:09:24.915360Z",
     "iopub.status.idle": "2022-04-06T23:13:16.310123Z",
     "shell.execute_reply": "2022-04-06T23:13:16.311433Z",
     "shell.execute_reply.started": "2022-04-06T00:30:21.722786Z"
    },
    "papermill": {
     "duration": 39831.436218,
     "end_time": "2022-04-06T23:13:16.311729",
     "exception": false,
     "start_time": "2022-04-06T12:09:24.875511",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87d4d2ce987a4961b76c714501c30b29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/52.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "11df971262d54988a02d032c14f4d542",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/474 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6bbb2ee159240d1a2cab1b856e849bb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/878k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e19b3c13a40140afb9574c5bbecdadfb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/446k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gradient accumulation 1\n",
      "Fold 0\n",
      "Epoch      Loss       lr  time\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "684aea155a224256a6f20c737a3c6661",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/533M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enable AWP\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:154: FutureWarning: Non-finite norm encountered in torch.nn.utils.clip_grad_norm_; continuing anyway. Note that the default behavior will change in a future release to error out if a non-finite total norm is encountered. At that point, setting error_if_nonfinite=false will be required to retain the old behavior.\n",
      "/opt/conda/lib/python3.7/site-packages/torch/optim/lr_scheduler.py:134: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
      "  \"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\", UserWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0.25 0.042260 0.020699 | 0.7415 1.99e-05   6.65 1.40 min\n",
      "Epoch 0.50 0.020987 0.016629 | 0.7800 1.95e-05  12.98 1.40 min\n",
      "Epoch 0.75 0.018189 0.013621 | 0.8110 1.89e-05  19.29 1.40 min\n",
      "Epoch 1.00 0.016576 0.015996 | 0.8244 1.81e-05  25.62 1.40 min\n",
      "Epoch 1.25 0.014337 0.014250 | 0.8491 1.71e-05  32.32 1.40 min\n",
      "Epoch 1.50 0.012059 0.012608 | 0.8545 1.59e-05  39.01 1.40 min\n",
      "Epoch 1.75 0.012621 0.012558 | 0.8561 1.45e-05  45.69 1.40 min\n",
      "Epoch 2.00 0.010895 0.013012 | 0.8582 1.31e-05  52.38 1.40 min\n",
      "Epoch 2.25 0.010602 0.012561 | 0.8618 1.16e-05  59.10 1.40 min\n",
      "Epoch 2.50 0.009421 0.013011 | 0.8598 1.00e-05  65.80 1.40 min\n",
      "Epoch 2.75 0.010040 0.013105 | 0.8657 8.44e-06  72.48 1.40 min\n",
      "Epoch 3.00 0.010766 0.012823 | 0.8652 6.91e-06  79.18 1.40 min\n",
      "Epoch 3.25 0.009007 0.013828 | 0.8669 5.46e-06  85.88 1.40 min\n",
      "Epoch 3.50 0.008715 0.013534 | 0.8679 4.12e-06  92.59 1.40 min\n",
      "Epoch 3.75 0.008269 0.013252 | 0.8657 2.93e-06  99.27 1.40 min\n",
      "Epoch 4.00 0.008151 0.013309 | 0.8688 1.91e-06 105.97 1.40 min\n",
      "Epoch 4.25 0.008230 0.014078 | 0.8683 1.09e-06 112.68 1.40 min\n",
      "Epoch 4.50 0.007380 0.014153 | 0.8690 4.89e-07 119.38 1.40 min\n",
      "Epoch 4.75 0.007055 0.014092 | 0.8687 1.23e-07 126.07 1.40 min\n",
      "Epoch 5.00 0.007508 0.014123 | 0.8695 0.00e+00 132.76 1.40 min\n",
      "Last model is same as best\n",
      "Fold 1\n",
      "Epoch      Loss       lr  time\n",
      "Enable AWP\n",
      "Epoch 0.25 0.043959 0.021489 | 0.7426 1.99e-05   6.36 1.41 min\n",
      "Epoch 0.50 0.020626 0.020100 | 0.7680 1.95e-05  12.68 1.41 min\n",
      "Epoch 0.75 0.017284 0.017150 | 0.8125 1.89e-05  19.00 1.41 min\n",
      "Epoch 1.00 0.015453 0.015681 | 0.8025 1.81e-05  25.33 1.41 min\n",
      "Epoch 1.25 0.013493 0.015593 | 0.8323 1.71e-05  32.03 1.41 min\n",
      "Epoch 1.50 0.012144 0.015488 | 0.8436 1.59e-05  38.71 1.41 min\n",
      "Epoch 1.75 0.010741 0.018022 | 0.8269 1.45e-05  45.40 1.41 min\n",
      "Epoch 2.00 0.011368 0.014739 | 0.8455 1.31e-05  52.09 1.41 min\n",
      "Epoch 2.25 0.009426 0.015487 | 0.8529 1.16e-05  58.80 1.41 min\n",
      "Epoch 2.50 0.010227 0.015332 | 0.8552 1.00e-05  65.51 1.41 min\n",
      "Epoch 2.75 0.009679 0.015159 | 0.8540 8.44e-06  72.20 1.41 min\n",
      "Epoch 3.00 0.008309 0.015837 | 0.8588 6.91e-06  78.91 1.41 min\n",
      "Epoch 3.25 0.008732 0.015956 | 0.8570 5.46e-06  85.63 1.41 min\n",
      "Epoch 3.50 0.009078 0.016372 | 0.8561 4.12e-06  92.33 1.41 min\n",
      "Epoch 3.75 0.008092 0.016301 | 0.8593 2.93e-06  99.03 1.41 min\n",
      "Epoch 4.00 0.007427 0.016015 | 0.8592 1.91e-06 105.74 1.41 min\n",
      "Epoch 4.25 0.007847 0.016632 | 0.8578 1.09e-06 112.46 1.41 min\n",
      "Epoch 4.50 0.007067 0.016390 | 0.8617 4.90e-07 119.15 1.41 min\n",
      "Epoch 4.75 0.008342 0.016476 | 0.8609 1.24e-07 125.84 1.41 min\n",
      "Epoch 5.00 0.007041 0.016502 | 0.8610 0.00e+00 132.56 1.41 min\n",
      "Last model is same as best\n",
      "Fold 2\n",
      "Epoch      Loss       lr  time\n",
      "Enable AWP\n",
      "Epoch 0.25 0.040404 0.020382 | 0.7185 1.99e-05   6.36 1.41 min\n",
      "Epoch 0.50 0.020786 0.016717 | 0.7921 1.95e-05  12.71 1.41 min\n",
      "Epoch 0.75 0.019738 0.013989 | 0.8219 1.89e-05  19.03 1.41 min\n",
      "Epoch 1.00 0.016069 0.017610 | 0.7942 1.81e-05  25.37 1.41 min\n",
      "Epoch 1.25 0.014107 0.014350 | 0.8511 1.71e-05  32.09 1.41 min\n",
      "Epoch 1.50 0.010555 0.014596 | 0.8520 1.59e-05  38.80 1.41 min\n",
      "Epoch 1.75 0.010577 0.013662 | 0.8543 1.45e-05  45.50 1.41 min\n",
      "Epoch 2.00 0.013219 0.013227 | 0.8570 1.31e-05  52.21 1.41 min\n",
      "Epoch 2.25 0.009385 0.015583 | 0.8576 1.16e-05  58.93 1.41 min\n",
      "Epoch 2.50 0.010712 0.014707 | 0.8593 1.00e-05  65.64 1.41 min\n",
      "Epoch 2.75 0.009828 0.013854 | 0.8617 8.44e-06  72.34 1.41 min\n",
      "Epoch 3.00 0.009258 0.013914 | 0.8690 6.91e-06  79.04 1.41 min\n",
      "Epoch 3.25 0.008273 0.014042 | 0.8681 5.46e-06  85.77 1.41 min\n",
      "Epoch 3.50 0.007711 0.014050 | 0.8684 4.12e-06  92.48 1.41 min\n",
      "Epoch 3.75 0.008109 0.014781 | 0.8657 2.93e-06  99.18 1.41 min\n",
      "Epoch 4.00 0.008805 0.014093 | 0.8663 1.91e-06 105.88 1.41 min\n",
      "Epoch 4.25 0.006613 0.015154 | 0.8688 1.09e-06 112.58 1.41 min\n",
      "Epoch 4.50 0.007481 0.014987 | 0.8685 4.89e-07 119.28 1.41 min\n",
      "Epoch 4.75 0.006444 0.015116 | 0.8683 1.23e-07 125.97 1.41 min\n",
      "Epoch 5.00 0.007831 0.015109 | 0.8695 0.00e+00 132.68 1.41 min\n",
      "Last model is same as best\n",
      "Fold 3\n",
      "Epoch      Loss       lr  time\n",
      "Enable AWP\n",
      "Epoch 0.25 0.046293 0.020846 | 0.7481 1.99e-05   6.36 1.42 min\n",
      "Epoch 0.50 0.021971 0.018070 | 0.7987 1.95e-05  12.69 1.42 min\n",
      "Epoch 0.75 0.018050 0.014765 | 0.8210 1.89e-05  19.02 1.42 min\n",
      "Epoch 1.00 0.016122 0.013585 | 0.8318 1.81e-05  25.35 1.42 min\n",
      "Epoch 1.25 0.012371 0.013369 | 0.8506 1.71e-05  32.04 1.42 min\n",
      "Epoch 1.50 0.011528 0.013703 | 0.8521 1.59e-05  38.74 1.42 min\n",
      "Epoch 1.75 0.011470 0.013831 | 0.8551 1.45e-05  45.44 1.42 min\n",
      "Epoch 2.00 0.010772 0.012637 | 0.8586 1.31e-05  52.12 1.42 min\n",
      "Epoch 2.25 0.010363 0.013339 | 0.8536 1.16e-05  58.82 1.42 min\n",
      "Epoch 2.50 0.010685 0.013227 | 0.8641 1.00e-05  65.50 1.42 min\n",
      "Epoch 2.75 0.009073 0.013280 | 0.8652 8.44e-06  72.19 1.42 min\n",
      "Epoch 3.00 0.010315 0.013139 | 0.8628 6.91e-06  78.87 1.42 min\n",
      "Epoch 3.25 0.008365 0.013453 | 0.8625 5.46e-06  85.58 1.42 min\n",
      "Epoch 3.50 0.008983 0.014284 | 0.8652 4.12e-06  92.27 1.42 min\n",
      "Epoch 3.75 0.009545 0.013006 | 0.8665 2.93e-06  98.97 1.42 min\n",
      "Epoch 4.00 0.008334 0.013570 | 0.8685 1.91e-06 105.66 1.42 min\n",
      "Epoch 4.25 0.007651 0.014089 | 0.8696 1.09e-06 112.35 1.42 min\n",
      "Epoch 4.50 0.008257 0.013906 | 0.8682 4.90e-07 119.06 1.42 min\n",
      "Epoch 4.75 0.008125 0.013934 | 0.8695 1.23e-07 125.78 1.42 min\n",
      "Epoch 5.00 0.007927 0.013920 | 0.8693 0.00e+00 132.50 1.42 min\n",
      "Last model is same as best\n",
      "Fold 4\n",
      "Epoch      Loss       lr  time\n",
      "Enable AWP\n",
      "Epoch 0.25 0.039692 0.020403 | 0.7449 1.99e-05   6.37 1.41 min\n",
      "Epoch 0.50 0.022107 0.016783 | 0.7961 1.95e-05  12.71 1.41 min\n",
      "Epoch 0.75 0.018007 0.015766 | 0.8094 1.89e-05  19.06 1.41 min\n",
      "Epoch 1.00 0.016156 0.013617 | 0.8333 1.81e-05  25.40 1.41 min\n",
      "Epoch 1.25 0.012760 0.013320 | 0.8521 1.71e-05  32.12 1.41 min\n",
      "Epoch 1.50 0.011621 0.013197 | 0.8555 1.59e-05  38.83 1.41 min\n",
      "Epoch 1.75 0.010889 0.013630 | 0.8605 1.45e-05  45.54 1.41 min\n",
      "Epoch 2.00 0.012085 0.012388 | 0.8590 1.31e-05  52.25 1.41 min\n",
      "Epoch 2.25 0.009550 0.012827 | 0.8658 1.16e-05  58.98 1.41 min\n",
      "Epoch 2.50 0.010530 0.013951 | 0.8674 1.00e-05  65.69 1.41 min\n",
      "Epoch 2.75 0.009206 0.014359 | 0.8631 8.44e-06  72.41 1.41 min\n",
      "Epoch 3.00 0.009418 0.014640 | 0.8575 6.91e-06  79.12 1.41 min\n",
      "Epoch 3.25 0.008698 0.014093 | 0.8671 5.46e-06  85.82 1.41 min\n",
      "Epoch 3.50 0.008131 0.014948 | 0.8662 4.12e-06  92.54 1.41 min\n",
      "Epoch 3.75 0.008414 0.014243 | 0.8719 2.93e-06  99.26 1.41 min\n",
      "Epoch 4.00 0.008963 0.013396 | 0.8726 1.91e-06 105.97 1.41 min\n",
      "Epoch 4.25 0.007347 0.014304 | 0.8708 1.09e-06 112.70 1.41 min\n",
      "Epoch 4.50 0.007405 0.014316 | 0.8719 4.90e-07 119.40 1.41 min\n",
      "Epoch 4.75 0.006870 0.014102 | 0.8724 1.23e-07 126.15 1.41 min\n",
      "Epoch 5.00 0.006838 0.014090 | 0.8728 0.00e+00 132.91 1.41 min\n",
      "Last model is same as best\n",
      "Train done 11.06 hr\n"
     ]
    }
   ],
   "source": [
    "debug = False\n",
    "tb_train = time.time()\n",
    "\n",
    "# Output directory\n",
    "name = 'deberta_base_awp'\n",
    "odir = name\n",
    "\n",
    "if not os.path.exists(odir):\n",
    "    os.mkdir(odir)\n",
    "\n",
    "# Data\n",
    "di = '/kaggle/input/nbme-score-clinical-patient-notes/'\n",
    "train = pd.read_csv(di + 'train.csv')    \n",
    "features = pd.read_csv(di + 'features.csv')\n",
    "patient_notes = pd.read_csv(di + 'patient_notes.csv')\n",
    "\n",
    "features['feature_text'] = features['feature_text'].apply(_replace_feature_text)\n",
    "\n",
    "# Attach text `pn_history` to train annotations\n",
    "train = train.merge(features, on=['feature_num', 'case_num'], how='left')\n",
    "train = train.merge(patient_notes, on=['pn_num', 'case_num'], how='left')\n",
    "\n",
    "nakama_fix_annotations(train)\n",
    "\n",
    "# Tokenizer\n",
    "transformer_path = 'microsoft/deberta-base'\n",
    "transformer_name = transformer_path.split('/')[-1]   # deberta-base\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(transformer_path)\n",
    "tokenizer.save_pretrained(odir + '/tokenizer')\n",
    "\n",
    "config = AutoConfig.from_pretrained(transformer_path, add_pooling_layer=False)\n",
    "config.save_pretrained(odir)\n",
    "\n",
    "# Tokenize data\n",
    "data = create_data(train, tokenizer)\n",
    "\n",
    "seed_everything(42)\n",
    "\n",
    "# Kfold\n",
    "nfold = 5\n",
    "folds = [0, 1, 2, 3, 4]\n",
    "\n",
    "kfold = StratifiedGroupKFold(nfold, shuffle=True, random_state=42)\n",
    "groups = train['pn_num'].values\n",
    "cases = train['case_num'].values\n",
    "\n",
    "# Parameters\n",
    "epochs = 5\n",
    "batch_size = 6\n",
    "batch_size_val = 16\n",
    "\n",
    "apex = True\n",
    "max_grad_norm = 1000\n",
    "optimizer_batch_size = 6  # batch size for optimization step\n",
    "gradient_accumulation_steps = max(optimizer_batch_size // batch_size, 1)\n",
    "print('gradient accumulation', gradient_accumulation_steps)\n",
    "\n",
    "eval_per_epoch = 4\n",
    "\n",
    "# Training and evaluation\n",
    "for ifold, (idx_train, idx_val) in enumerate(kfold.split(data, cases, groups=groups)):\n",
    "    if ifold not in folds:\n",
    "        continue\n",
    "\n",
    "    tb = time.time()\n",
    "    print('Fold %d' % ifold)\n",
    "    print('Epoch      Loss       lr  time')\n",
    "\n",
    "    # Train - validation split\n",
    "    data_train = data[idx_train]\n",
    "    data_val = data[idx_val]\n",
    "\n",
    "    if debug:\n",
    "        data_train = data_train[:64]\n",
    "        data_val = data_val[:16]\n",
    "        epochs = 2\n",
    "        folds = [0]\n",
    "\n",
    "    loader_train = DataLoader(Dataset(data_train),\n",
    "                              batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "    loader_val = DataLoader(Dataset(data_val), batch_size=batch_size_val)\n",
    "\n",
    "    # Model\n",
    "    model = Model(transformer_path)\n",
    "    model.train()\n",
    "    model.to(device)\n",
    "\n",
    "    # Optimizer\n",
    "    lr = 2e-5\n",
    "    eps = 1e-6\n",
    "    betas = (0.9, 0.999)\n",
    "    weight_decay = 0.01\n",
    "    optimizer_parameters = get_optimizer_params(model,\n",
    "                                                encoder_lr=lr,\n",
    "                                                decoder_lr=lr,\n",
    "                                                weight_decay=weight_decay)\n",
    "\n",
    "    optimizer = AdamW(optimizer_parameters, lr=lr, eps=eps, betas=betas)\n",
    "\n",
    "    # One `step` is one optimizer step, `gradient_accumulation` mini batches\n",
    "    steps_per_epoch = len(loader_train) // gradient_accumulation_steps\n",
    "    eval_steps = [int(i * steps_per_epoch / eval_per_epoch) for i in range(1, eval_per_epoch)] + \\\n",
    "                 [steps_per_epoch]\n",
    "\n",
    "    # Scheduler\n",
    "    num_train_steps = steps_per_epoch * epochs\n",
    "    num_warmup_steps = 0\n",
    "    num_cycles = 0.5\n",
    "    scheduler = get_cosine_schedule_with_warmup(optimizer,\n",
    "                                                num_warmup_steps=num_warmup_steps,\n",
    "                                                num_training_steps=num_train_steps,\n",
    "                                                num_cycles=num_cycles)\n",
    "\n",
    "    criterion = nn.BCEWithLogitsLoss(reduction='none')\n",
    "    best_score = 0\n",
    "    best_epoch = None\n",
    "\n",
    "    print('Enable AWP')\n",
    "    awp = AWP(model, optimizer, adv_lr=0.001, adv_eps=0.001)\n",
    "    awp_start = 1.0\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        scaler = torch.cuda.amp.GradScaler(enabled=apex)\n",
    "\n",
    "        step = 0\n",
    "        n_sum = 0\n",
    "        loss_sum = 0\n",
    "        for ibatch, d in enumerate(loader_train):\n",
    "            input_ids = d['input_ids'].to(device)\n",
    "            attention_mask = d['attention_mask'].to(device)\n",
    "            y = d['y'].to(device)\n",
    "            n = y.size(0)\n",
    "\n",
    "            if epoch >= awp_start:\n",
    "                awp.perturb(input_ids, attention_mask, y, criterion)\n",
    "\n",
    "            with torch.cuda.amp.autocast(enabled=apex):\n",
    "                y_pred = model(input_ids, attention_mask)\n",
    "\n",
    "            loss = criterion(y_pred.view(-1), y.view(-1))\n",
    "            loss = torch.masked_select(loss, y.view(-1) != -1).mean()\n",
    "            loss_sum += n * loss.item()\n",
    "            n_sum += n\n",
    "\n",
    "            if gradient_accumulation_steps > 1:\n",
    "                loss = loss / gradient_accumulation_steps\n",
    "\n",
    "            scaler.scale(loss).backward()\n",
    "            awp.restore()\n",
    "\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n",
    "\n",
    "            if (ibatch + 1) % gradient_accumulation_steps == 0:\n",
    "                scaler.step(optimizer)\n",
    "                scaler.update()\n",
    "                optimizer.zero_grad()\n",
    "                if scheduler is not None:\n",
    "                    scheduler.step()\n",
    "                step += 1\n",
    "\n",
    "            del loss, y_pred, input_ids, attention_mask, y\n",
    "\n",
    "            ep = epoch + step / steps_per_epoch\n",
    "\n",
    "            # Validation\n",
    "            if step in eval_steps and (ibatch + 1) % gradient_accumulation_steps == 0:\n",
    "                val = evaluate(model, loader_val, criterion)\n",
    "                score = compute_score(val['preds'], data_val)\n",
    "\n",
    "                ep = epoch + step / steps_per_epoch\n",
    "                loss_train = loss_sum / n_sum\n",
    "                lr1 = optimizer.param_groups[0]['lr']\n",
    "                dt = (time.time() - tb) / 60\n",
    "\n",
    "                print('Epoch %.2f %.6f %.6f | %.4f %.2e %6.2f %.2f min' %\n",
    "                      (ep, loss_train, val['loss'], score,\n",
    "                       lr1, dt, val['time'] / 60))\n",
    "\n",
    "                n_sum = loss_sum = 0\n",
    "\n",
    "                if step == steps_per_epoch:\n",
    "                    break  # drop incomplete gradient accumulation\n",
    "\n",
    "        # Save model\n",
    "        if score >= best_score:\n",
    "            best_score = score\n",
    "            best_epoch = epoch + 1\n",
    "\n",
    "            model_filename = 'model%d_best.pytorch' % ifold\n",
    "            torch.save(model.state_dict(), model_filename)\n",
    "            #print(model_filename, 'written')\n",
    "\n",
    "    # Save final model\n",
    "    if best_epoch == epochs:\n",
    "        print('Last model is same as best')\n",
    "    #else:\n",
    "        #model_filename = '%s/model%d.pytorch' % (odir, ifold)\n",
    "        #model.to('cpu')\n",
    "        #model.eval()\n",
    "        #torch.save(model.state_dict(), model_filename)\n",
    "        #print(model_filename, 'written')\n",
    "\n",
    "    del model\n",
    "    del awp\n",
    "\n",
    "dt = (time.time() - tb_train) / 3600\n",
    "print('Train done %.2f hr' % dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c4e7d3b3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-06T23:13:16.449879Z",
     "iopub.status.busy": "2022-04-06T23:13:16.448882Z",
     "iopub.status.idle": "2022-04-06T23:13:18.108797Z",
     "shell.execute_reply": "2022-04-06T23:13:18.108190Z",
     "shell.execute_reply.started": "2022-04-06T00:31:04.397307Z"
    },
    "papermill": {
     "duration": 1.719713,
     "end_time": "2022-04-06T23:13:18.108979",
     "exception": false,
     "start_time": "2022-04-06T23:13:16.389266",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__notebook__.ipynb   model1_best.pytorch  model4_best.pytorch\r\n",
      "deberta_base_awp     model2_best.pytorch\r\n",
      "model0_best.pytorch  model3_best.pytorch\r\n",
      "config.json  tokenizer\r\n"
     ]
    }
   ],
   "source": [
    "! ls\n",
    "! ls deberta_base_awp"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.7.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 39851.345708,
   "end_time": "2022-04-06T23:13:21.207619",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2022-04-06T12:09:09.861911",
   "version": "2.3.3"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "00e439af9e3244709ecc090494cbab15": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "0b520f6cb32045fd832f182ba1c09a81": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "0e3f277f8be54b21b4ec659074f5f5b8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "0f9b7cc6a6df4e959ab69ff12d39e941": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_74d2cbb0655f44ef8be4f3494d57e909",
       "max": 558614189.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_62152ed7442a4bca8837108f3374ce56",
       "value": 558614189.0
      }
     },
     "10858dd36168419d9fc4a98d0f9c3cc8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "11df971262d54988a02d032c14f4d542": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_15cf69a2a57e483f86d6f0f6ff3bc1f8",
        "IPY_MODEL_2af27cf9b472414c8485d91804a78c55",
        "IPY_MODEL_e144b808ba6c49f19ecb57af338d4aa5"
       ],
       "layout": "IPY_MODEL_abd4c1205d04416b98f4eafc3785b01f"
      }
     },
     "11f1728824a94c61abfdeb9c7956821f": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "15cf69a2a57e483f86d6f0f6ff3bc1f8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_39a616ab6d6c4d218841f7401cf65f06",
       "placeholder": "​",
       "style": "IPY_MODEL_1e610f7d9970474db325e6b261fbb6eb",
       "value": "Downloading: 100%"
      }
     },
     "1b03956b5bdf407793873b8e7cb295f1": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "1e610f7d9970474db325e6b261fbb6eb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "2879b45b71c441438d81d301d682396f": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_abe2b0c7a3184a0384ffb4dc05a2b569",
       "placeholder": "​",
       "style": "IPY_MODEL_f6be667d161243e0b438844d252140ce",
       "value": " 446k/446k [00:00&lt;00:00, 4.53MB/s]"
      }
     },
     "2af27cf9b472414c8485d91804a78c55": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_890f1bcf9db84ea79566a1ddd544a092",
       "max": 474.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_8e6f7c7e80074436a8267f43cb15588c",
       "value": 474.0
      }
     },
     "2b6d763f90b74907bc4d3bf33dce533a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "34e6d3da2d1e45dbbd6ecf19436a72ab": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "37fb6a6897ff45108c0f181844ad864a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_d8ac7232d1454c6b8d568bc3f3aa3ce0",
       "max": 52.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_6737705adc2f406491aff1ffadf9fd52",
       "value": 52.0
      }
     },
     "39a616ab6d6c4d218841f7401cf65f06": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "3f29e402fea949448f50a979079b1843": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "4c4fb3fdfb9640799e268af9783d4bd1": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_1b03956b5bdf407793873b8e7cb295f1",
       "placeholder": "​",
       "style": "IPY_MODEL_e01764bc5850446e9dc9fa4f192473f0",
       "value": " 52.0/52.0 [00:00&lt;00:00, 2.01kB/s]"
      }
     },
     "55c9b6c1713b4929943e67dbbcbae640": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "62152ed7442a4bca8837108f3374ce56": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "6737705adc2f406491aff1ffadf9fd52": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "684aea155a224256a6f20c737a3c6661": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_ef0c353558744ff29cee5bd915590abc",
        "IPY_MODEL_0f9b7cc6a6df4e959ab69ff12d39e941",
        "IPY_MODEL_727c6578df0444d69f2a6fdc6f0b03d5"
       ],
       "layout": "IPY_MODEL_00e439af9e3244709ecc090494cbab15"
      }
     },
     "68b3311308f546f8bdf5b42851612be2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "6f3e05abfe614aedafa13b6368d55417": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "70dbddcae192459d9a28b65a6c0212de": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "727c6578df0444d69f2a6fdc6f0b03d5": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_34e6d3da2d1e45dbbd6ecf19436a72ab",
       "placeholder": "​",
       "style": "IPY_MODEL_faa1b016113e40b48fd201e7cd1909ae",
       "value": " 533M/533M [00:11&lt;00:00, 50.0MB/s]"
      }
     },
     "74d2cbb0655f44ef8be4f3494d57e909": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "84b3d53200b1468986c761753451032d": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "87d4d2ce987a4961b76c714501c30b29": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_8fb4a8236a684928aadaef7dc831583c",
        "IPY_MODEL_37fb6a6897ff45108c0f181844ad864a",
        "IPY_MODEL_4c4fb3fdfb9640799e268af9783d4bd1"
       ],
       "layout": "IPY_MODEL_70dbddcae192459d9a28b65a6c0212de"
      }
     },
     "880315da3cdb41019ada14bde6a4de60": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "890f1bcf9db84ea79566a1ddd544a092": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "8e6f7c7e80074436a8267f43cb15588c": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "8fb4a8236a684928aadaef7dc831583c": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_0b520f6cb32045fd832f182ba1c09a81",
       "placeholder": "​",
       "style": "IPY_MODEL_10858dd36168419d9fc4a98d0f9c3cc8",
       "value": "Downloading: 100%"
      }
     },
     "92c1f743bbc1433eb23138f6412ac4df": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "92f8e5af03e442b5827be5d07b68527a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "94ef669bc36c4fa5a42c603cf200ca3a": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "9ac4c4c2c0e14e0c8f97aed833c9b42c": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_94ef669bc36c4fa5a42c603cf200ca3a",
       "placeholder": "​",
       "style": "IPY_MODEL_55c9b6c1713b4929943e67dbbcbae640",
       "value": "Downloading: 100%"
      }
     },
     "9f052537f0f64ad197a5cad950d1ac24": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "9f606ac1b85842c28dc7dba7a10a8b9e": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_e24f3947da704622969fd68f87ed66a5",
       "placeholder": "​",
       "style": "IPY_MODEL_0e3f277f8be54b21b4ec659074f5f5b8",
       "value": " 878k/878k [00:00&lt;00:00, 5.32MB/s]"
      }
     },
     "abd4c1205d04416b98f4eafc3785b01f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "abe2b0c7a3184a0384ffb4dc05a2b569": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "c6baeaf7172e4d84871c2004fdae9eb7": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "cc005d77d9f5449887b783f49049d3a4": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_68b3311308f546f8bdf5b42851612be2",
       "max": 898825.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_92c1f743bbc1433eb23138f6412ac4df",
       "value": 898825.0
      }
     },
     "d2c57bde5fa2460783df86a186cfede7": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_84b3d53200b1468986c761753451032d",
       "max": 456318.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_11f1728824a94c61abfdeb9c7956821f",
       "value": 456318.0
      }
     },
     "d780053f8b1e44f3a10d73f915159521": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "d8ac7232d1454c6b8d568bc3f3aa3ce0": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "e01764bc5850446e9dc9fa4f192473f0": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "e144b808ba6c49f19ecb57af338d4aa5": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_6f3e05abfe614aedafa13b6368d55417",
       "placeholder": "​",
       "style": "IPY_MODEL_d780053f8b1e44f3a10d73f915159521",
       "value": " 474/474 [00:00&lt;00:00, 19.5kB/s]"
      }
     },
     "e19b3c13a40140afb9574c5bbecdadfb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_9ac4c4c2c0e14e0c8f97aed833c9b42c",
        "IPY_MODEL_d2c57bde5fa2460783df86a186cfede7",
        "IPY_MODEL_2879b45b71c441438d81d301d682396f"
       ],
       "layout": "IPY_MODEL_3f29e402fea949448f50a979079b1843"
      }
     },
     "e24f3947da704622969fd68f87ed66a5": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "e6bbb2ee159240d1a2cab1b856e849bb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_f1b206a0e924400baaea9e16bbd03829",
        "IPY_MODEL_cc005d77d9f5449887b783f49049d3a4",
        "IPY_MODEL_9f606ac1b85842c28dc7dba7a10a8b9e"
       ],
       "layout": "IPY_MODEL_880315da3cdb41019ada14bde6a4de60"
      }
     },
     "ef0c353558744ff29cee5bd915590abc": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_9f052537f0f64ad197a5cad950d1ac24",
       "placeholder": "​",
       "style": "IPY_MODEL_2b6d763f90b74907bc4d3bf33dce533a",
       "value": "Downloading: 100%"
      }
     },
     "f1b206a0e924400baaea9e16bbd03829": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_c6baeaf7172e4d84871c2004fdae9eb7",
       "placeholder": "​",
       "style": "IPY_MODEL_92f8e5af03e442b5827be5d07b68527a",
       "value": "Downloading: 100%"
      }
     },
     "f6be667d161243e0b438844d252140ce": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "faa1b016113e40b48fd201e7cd1909ae": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}