Finetune_GLUE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "ugFdx5v_YxzA",
    "outputId": "e6663d66-4c4b-4953-8388-7fa2b23b126a",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os,sys\n",
    "from pathlib import Path\n",
    "from functools import partial\n",
    "from inspect import isclass\n",
    "import random\n",
    "from IPython.core.debugger import set_trace as bk\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import datasets\n",
    "from fastai.text.all import *\n",
    "from transformers import ElectraModel, ElectraConfig, ElectraTokenizerFast, ElectraForPreTraining\n",
    "from hugdatafast import *\n",
    "from _utils.utils import *\n",
    "from _utils.would_like_to_pr import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Confiquration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "c = MyConfig({\n",
    "\n",
    "  'device': 'cuda:0', #List[int]: use multi gpu (data parallel)\n",
    "  # run [start,end) runs, every run finetune every GLUE tasks once with different seeds.\n",
    "  'start':0,\n",
    "  'end': 10,\n",
    "  \n",
    "  'pretrained_checkpoint': 'vanilla_11081_100.0%.pth', # None to use pretrained ++ model from HuggingFace\n",
    "  # Seeds for fintuning. i th run use i th seeds, None to use system time\n",
    "  #'seeds': [939, 481, 569, 620, 159, 808, 816, 101, 554, 104], # for 11081\n",
    "  #'seeds': [611, 609, 830, 237, 668, 608, 475, 690, 53, 94], # for 36\n",
    "  #'seeds': [775, 961, 778, 915, 979, 526, 99, 669, 806, 78], # for 1188\n",
    "  #'seeds': [895, 602, 573, 457, 736, 871, 571, 84, 514, 740,], # for 76\n",
    "  #'seeds': [760, 63, 392, 240, 794, 168, 245, 345, 97, 917], # 1\n",
    "  #'seeds': [6669, 4093, 6254, 8546, 489, 901, 5567, 3690, 7057, 3663,], # for 4\n",
    "  #'seeds': [3426, 2730, 6509, 6957, 2961, 7783, 7061, 4261, 2256, 4863,], # for 4649\n",
    "  #'seeds': [3049, 3005, 3298, 8108, 6676, 2275, 376, 5053, 154, 8992,], # for 7\n",
    "  'seeds': None,\n",
    "\n",
    "  'weight_decay': 0,\n",
    "  'adam_bias_correction': False,\n",
    "  'xavier_reinited_outlayer': True,\n",
    "  'schedule': 'original_linear',\n",
    "  'original_lr_layer_decays': True,\n",
    "  'double_unordered': True,\n",
    "  \n",
    "  # whether to do finetune or test\n",
    "  'do_finetune': True, # True -> do finetune ; False -> do test\n",
    "  # finetuning checkpoint for testing. These will become \"ckp_dir/{task}_{group_name}_{th_run}.pth\"\n",
    "  'th_run': { 'qqp': 7, 'qnli': 5,\n",
    "              'mrpc': 7, 'mnli': 2, 'ax': 2,\n",
    "              'sst2': 3, 'rte': 7,  'wnli': 0, \n",
    "              'cola': 1, 'stsb': 8,  \n",
    "            },\n",
    "  \n",
    "  'size': 'small',\n",
    "  'wsc_trick': False,\n",
    "\n",
    "  'num_workers': 3,\n",
    "  'my_model': False, # True only for my personal research\n",
    "  'logger': 'wandb',\n",
    "  'group_name': None, # the name of represents these runs\n",
    "  # None: use name of checkpoint.\n",
    "  # False: don't do online logging and don't save checkpoints\n",
    "})\n",
    "\n",
    "# only for my personal research purpose\n",
    "hparam_update = {\n",
    "  \n",
    "}\n",
    "\n",
    "\"\"\" Vanilla ELECTRA settings\n",
    "'adam_bias_correction': False,\n",
    "'xavier_reinited_outlayer': True,\n",
    "'schedule': 'original_linear',\n",
    "'original_lr_layer_decays': True,\n",
    "'double_unordered': True,\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Check\n",
    "if not c.do_finetune: assert c.th_run['mnli'] == c.th_run['ax']\n",
    "if c.pretrained_checkpoint is None: assert not c.my_model\n",
    "assert c.schedule in ['original_linear', 'separate_linear', 'one_cycle', 'adjusted_one_cycle']\n",
    "\n",
    "# Settings of different sizes\n",
    "if c.size == 'small': c.lr = 3e-4; c.layer_lr_decay = 0.8; c.max_length = 128\n",
    "elif c.size == 'base': c.lr = 1e-4; c.layer_lr_decay = 0.8; c.max_length = 512\n",
    "elif c.size == 'large': c.lr = 5e-5; c.layer_lr_decay = 0.9; c.max_length = 512\n",
    "else: raise ValueError(f\"Invalid size {c.size}\")\n",
    "if c.pretrained_checkpoint is None: c.max_length = 512 # All public models is ++, which use max_length 512\n",
    "\n",
    "# huggingface/transformers\n",
    "hf_tokenizer = ElectraTokenizerFast.from_pretrained(f\"google/electra-{c.size}-discriminator\")\n",
    "electra_config = ElectraConfig.from_pretrained(f'google/electra-{c.size}-discriminator')\n",
    "\n",
    "# wsc\n",
    "if c.wsc_trick:\n",
    "  from _utils.wsc_trick import * # importing spacy model takes time\n",
    "\n",
    "# logging\n",
    "# light logging callback here is to only log the last score and avoid exceeding the api access limit\n",
    "if c.logger == 'neptune':\n",
    "  import neptune\n",
    "  from fastai.callback.neptune import NeptuneCallback\n",
    "  class LightNeptuneCallback(NeptuneCallback):\n",
    "    def after_batch(self): pass\n",
    "    def after_epoch(self):\n",
    "      if self.epoch == (self.n_epoch - 1): super().after_epoch()\n",
    "  neptune.init(project_qualified_name='richard-wang/electra-glue')\n",
    "elif c.logger == 'wandb':\n",
    "  import wandb\n",
    "  from fastai.callback.wandb import WandbCallback\n",
    "  class LightWandbCallback(Callback):\n",
    "    def __init__(self, run):\n",
    "      self.run = run\n",
    "    def after_epoch(self):\n",
    "      if self.epoch != (self.n_epoch - 1): return\n",
    "      wandb.log({n:s for n,s in zip(self.recorder.metric_names, self.recorder.log) if n not in ['train_loss', 'epoch', 'time']})\n",
    "    def after_fit(self):\n",
    "      wandb.log({}) # ensure sync of last step\n",
    "      self.run.finish()\n",
    "\n",
    "# my model\n",
    "if c.my_model:\n",
    "  sys.path.insert(0, os.path.abspath(\"..\"))\n",
    "  from modeling.model import ModelForDiscriminator\n",
    "  from hyperparameter import electra_hparam_from_hf\n",
    "  hparam = electra_hparam_from_hf(electra_config, hf_tokenizer)\n",
    "  hparam.update(hparam_update)\n",
    "\n",
    "# Path\n",
    "Path('./datasets').mkdir(exist_ok=True)\n",
    "Path('./checkpoints/glue').mkdir(exist_ok=True, parents=True)\n",
    "Path('./test_outputs/glue').mkdir(exist_ok=True, parents=True)\n",
    "c.pretrained_ckp_path = Path(f'./checkpoints/pretrain/{c.pretrained_checkpoint}')\n",
    "if c.group_name is None:\n",
    "  if c.pretrained_checkpoint: c.group_name = c.pretrained_checkpoint[:-4]\n",
    "  elif c.pretrained_checkpoint is None: c.group_name = f\"{c.size}++\"\n",
    "\n",
    "# Print info\n",
    "print(f\"process id: {os.getpid()}\")\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "METRICS = {\n",
    "  **{ task:[MatthewsCorrCoef()] for task in ['cola']},\n",
    "  **{ task:[accuracy] for task in ['sst2', 'mnli', 'qnli', 'rte', 'wnli', 'snli','ax']},\n",
    "  **{ task:[F1Score(), accuracy] for task in ['mrpc', 'qqp']}, \n",
    "  **{ task:[PearsonCorrCoef(), SpearmanCorrCoef()] for task in ['stsb']}\n",
    "}\n",
    "NUM_CLASS = {\n",
    "    **{ task:1 for task in ['stsb']},\n",
    "    **{ task:2 for task in ['cola', 'sst2', 'mrpc', 'qqp', 'qnli', 'rte', 'wnli']},\n",
    "    **{ task:3 for task in ['mnli','ax']},\n",
    "}\n",
    "TEXT_COLS = {\n",
    "    **{ task:['question', 'sentence'] for task in ['qnli']},\n",
    "    **{ task:['sentence1', 'sentence2'] for task in ['mrpc','stsb','wnli','rte']},\n",
    "    **{ task:['question1','question2'] for task in ['qqp']},\n",
    "    **{ task:['premise','hypothesis'] for task in ['mnli','ax']},\n",
    "    **{ task:['sentence'] for task in ['cola','sst2']},\n",
    "}\n",
    "LOSS_FUNC = {\n",
    "    **{ task: CrossEntropyLossFlat() for task in ['cola','sst2','mrpc','qqp','mnli','qnli','rte','wnli', 'ax']},\n",
    "    **{ task: MyMSELossFlat(low=0.0, high=5.0) for task in ['stsb']}\n",
    "}\n",
    "if c.wsc_trick: \n",
    "  LOSS_FUNC['wnli'] = ELECTRAWSCTrickLoss()\n",
    "  METRICS['wnli'] = [wsc_trick_accuracy]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yoI6QffPY6Jv"
   },
   "source": [
    "# 2. Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "p_P8P4hffUr8"
   },
   "source": [
    "## 2.1 Download and Preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_sents_max_len(example, cols, max_len, swap=False):\n",
    "  # Follow BERT and ELECTRA, truncate the examples longer than max length\n",
    "  tokens_a = hf_tokenizer.tokenize(example[cols[0]])\n",
    "  tokens_b = hf_tokenizer.tokenize(example[cols[1]]) if len(cols)==2 else []\n",
    "  _max_length = max_len - 1 - len(cols) # preserved for cls and sep tokens\n",
    "  while True:\n",
    "    total_length = len(tokens_a) + len(tokens_b)\n",
    "    if total_length <= _max_length:\n",
    "      break\n",
    "    if len(tokens_a) > len(tokens_b):\n",
    "      tokens_a.pop()\n",
    "    else:\n",
    "      tokens_b.pop()\n",
    "  if swap:\n",
    "    tokens_a, tokens_b = tokens_b, tokens_a\n",
    "  tokens = [hf_tokenizer.cls_token, *tokens_a, hf_tokenizer.sep_token]\n",
    "  token_type = [0]*len(tokens)\n",
    "  if tokens_b: \n",
    "    tokens += [*tokens_b, hf_tokenizer.sep_token]\n",
    "    token_type += [1]*(len(tokens_b)+1)\n",
    "  example['inp_ids'] = hf_tokenizer.convert_tokens_to_ids(tokens)\n",
    "  example['attn_mask'] = [1] * len(tokens)\n",
    "  example['token_type_ids'] = token_type\n",
    "  return example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "referenced_widgets": [
      "ae0863b074ff4a1a959514881fe2802f",
      "e5d5d86746ec4e2298028c57ef775d84",
      "a8b32129b9ab41d0895cb1ec445db9ca",
      "2d26c89a3d97493e911c52fa20cf79d9",
      "65eee3543a9d46adbc8037995c3a69a0",
      "9fc29a98d6be46e0974d1f0f7681d5c3",
      "5b8f1e5172204ed58bd60972289b24a8",
      "bd8da3c0b2784f2094d3b0cfa744cf87",
      "f74af04bf4764b7888f0a45e37ce85c1",
      "f35afd9d13c240aba4e0619cca0199b7",
      "a07124360f4c4249a745891d83e719f6",
      "f26a4f6d9334465d9022ab2a3401035b",
      "007556af2af8413ca47aa5f7f725e68a",
      "8d4771f53c7046f1afe7e04a3a647a31",
      "71beed6de71c4aa2862ef33e21cc58bc",
      "e1d8bdd9c1a74bc18a103bc900589c2c"
     ]
    },
    "colab_type": "code",
    "id": "U-PLcDIRY6y-",
    "outputId": "12e31f13-b70f-48c8-dd45-c6f1718b03c1",
    "tags": []
   },
   "outputs": [],
   "source": [
    "glue_dsets = {}; glue_dls = {}\n",
    "for task in ['cola', 'sst2', 'mrpc', 'stsb', 'mnli', 'qqp', 'qnli', 'rte', 'wnli', 'ax']:\n",
    "\n",
    "  # Load / download datasets.\n",
    "  dsets = datasets.load_dataset('glue', task, cache_dir='./datasets')\n",
    "\n",
    "  # There is two samples broken in QQP training set\n",
    "  if task=='qqp': dsets['train'] = dsets['train'].filter(lambda e: e['question2']!='',\n",
    "                        cache_file_name=os.path.join(dsets['train'].cache_directory(), 'fixed_train.arrow'))\n",
    "\n",
    "  # Load / Make tokenized datasets\n",
    "  tok_func = partial(tokenize_sents_max_len, cols=TEXT_COLS[task], max_len=c.max_length)\n",
    "  glue_dsets[task] = dsets.my_map(tok_func, cache_file_names=f\"tokenized_{c.max_length}_{{split}}\")\n",
    "\n",
    "  if c.double_unordered and task in ['mrpc', 'stsb']:\n",
    "    swap_tok_func = partial(tokenize_sents_max_len, cols=TEXT_COLS[task], max_len=c.max_length, swap=True)\n",
    "    swapped_train = dsets['train'].my_map(swap_tok_func, \n",
    "                                          cache_file_name=f\"swapped_tokenized_{c.max_length}_train\")\n",
    "    glue_dsets[task]['train'] = datasets.concatenate_datasets([glue_dsets[task]['train'], swapped_train])\n",
    "\n",
    "  # Load / Make dataloaders\n",
    "  hf_dsets = HF_Datasets(glue_dsets[task], hf_toker=hf_tokenizer, n_inp=3,\n",
    "                cols={'inp_ids':TensorText, 'attn_mask':noop, 'token_type_ids':noop, 'label':TensorCategory})\n",
    "  if c.double_unordered and task in ['mrpc', 'stsb']:\n",
    "    dl_kwargs = {'train': {'cache_name': f\"double_dl_{c.max_length}_train.json\"}}\n",
    "  else: dl_kwargs = None\n",
    "  glue_dls[task] = hf_dsets.dataloaders(bs=32, shuffle_train=True, num_workers=c.num_workers,\n",
    "                                        cache_name=f\"dl_{c.max_length}_{{split}}.json\",\n",
    "                                        dl_kwargs=dl_kwargs,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if c.wsc_trick:\n",
    "  wsc = datasets.load_dataset('super_glue', 'wsc', cache_dir='./datasets')\n",
    "  glue_dsets['wnli'] = wsc.my_map(partial(wsc_trick_process, hf_toker=hf_tokenizer),\n",
    "                                  cache_file_names=\"tricked_{split}.arrow\")\n",
    "  cols={'prefix':TensorText,'suffix':TensorText,'cands':TensorText,'cand_lens':noop,'label':TensorCategory}\n",
    "  glue_dls['wnli'] = HF_Datasets(glue_dsets['wnli'], hf_toker=hf_tokenizer, n_inp=4, \n",
    "                                 cols=cols).dataloaders(bs=32, cache_name=\"dl_tricked_{split}.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MqtDnxK9ZYQP"
   },
   "source": [
    "## 1.2 View Data\n",
    "- View raw data on [nlp-viewer]! (https://huggingface.co/nlp/viewer/)\n",
    "\n",
    "- View task description on Tensorflow dataset doc for GLUE (https://www.tensorflow.org/datasets/catalog/glue) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if False:\n",
    "  print(\"CoLA (The Corpus of Linguistic Acceptability) - 0: unacceptable, 1: acceptable\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['cola'].loaders]))\n",
    "  glue_dls['cola'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"SST-2 (The Stanford Sentiment Treebank) - 1: positvie, 0: negative\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['sst2'].loaders]))\n",
    "  glue_dls['sst2'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"MRPC (Microsoft Research Paraphrase Corpus) -  1: match, 0: no\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['mrpc'].loaders]))\n",
    "  glue_dls['mrpc'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"STS-B (Semantic Textual Similarity Benchmark) - 0.0 ~ 5.0\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['stsb'].loaders]))\n",
    "  glue_dls['stsb'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"QQP (Quora Question Pairs) - 0: no, 1: duplicated\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['qqp'].loaders]))\n",
    "  glue_dls['qqp'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"MNLI (The Multi-Genre NLI Corpus) - 0: entailment, 1: neutral, 2: contradiction\")\n",
    "  print(\"Dataset size (train/validation_matched/validation_mismatched/test_matched/test_mismatched): {}/{}/{}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['mnli'].loaders]))\n",
    "  glue_dls['mnli'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"(QNLI (The Stanford Question Answering Dataset) - 0: entailment, 1: not_entailment)\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['qnli'].loaders]))\n",
    "  glue_dls['qnli'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"RTE (Recognizing_Textual_Entailment) - 0: entailment, 1: not_entailment\")\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['rte'].loaders]))\n",
    "  glue_dls['rte'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"WSC (The Winograd Schema Challenge) - 0: wrong, 1: correct\")\n",
    "  # There are three style, WNLI (casted in NLI type), WSC, WSC with candidates (trick used by Roberta)\n",
    "  \"Note for WSC trick: cands is the concatenation of candidates, cand_lens is the lengths of candidates in order.\"\n",
    "  print(\"Dataset size (train/valid/test): {}/{}/{}\".format(*[len(dl.dataset) for dl in glue_dls['wnli'].loaders]))\n",
    "  glue_dls['wnli'].show_batch(max_n=1)\n",
    "  print()\n",
    "  print(\"AX (GLUE Diagnostic Dataset) - 0: entailment, 1: neutral, 2: contradiction\")\n",
    "  print(\"Dataset size (test): {}\".format(*[len(dl.dataset) for dl in glue_dls['ax'].loaders]))\n",
    "  glue_dls['ax'].show_batch(max_n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zi5_OKHqX8qw"
   },
   "source": [
    "# 2. Finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Finetuning model\n",
    "* ELECTRA use CLS encodings as pooled result to predict the sentence. (see [here](https://github.com/google-research/electra/blob/79111328070e491b287c307906701ebc61091eb2/model/modeling.py#L254) of its official repository)\n",
    "\n",
    "* Note that we should use different prediction head instance for different tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MF6vajgKQKJ7"
   },
   "outputs": [],
   "source": [
    "class SentencePredictor(nn.Module):\n",
    "\n",
    "  def __init__(self, model, hidden_size, num_class):\n",
    "    super().__init__()\n",
    "    self.base_model = model\n",
    "    self.dropout = nn.Dropout(0.1)\n",
    "    self.classifier = nn.Linear(hidden_size, num_class)\n",
    "    if c.xavier_reinited_outlayer:\n",
    "      nn.init.xavier_uniform_(self.classifier.weight.data)\n",
    "      self.classifier.bias.data.zero_()\n",
    "\n",
    "  def forward(self, input_ids, attention_mask, token_type_ids):\n",
    "    x = self.base_model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)[0]\n",
    "    return self.classifier(self.dropout(x[:,0,:])).squeeze(-1).float() # if regression task, squeeze to (B), else (B,#class)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "poNJYRwETtBo"
   },
   "source": [
    "## 2.2 Discriminative learning rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7OyT-hy-7FXz"
   },
   "outputs": [],
   "source": [
    "def list_parameters(model, submod_name):\n",
    "  return list(eval(f\"model.{submod_name}\").parameters())\n",
    "\n",
    "def hf_electra_param_splitter(model, wsc_trick=False):\n",
    "  base = 'discriminator.electra' if wsc_trick else 'base_model'\n",
    "  embed_name = 'embedding' if c.my_model else 'embeddings'\n",
    "  scaler_name = 'dimension_scaler' if c.my_model else 'embeddings_project'\n",
    "  layers_name = 'layers' if c.my_model else 'layer'\n",
    "  output_name = 'classifier' if not wsc_trick else f'discriminator.discriminator_predictions'\n",
    "  \n",
    "  groups = [ list_parameters(model, f\"{base}.{embed_name}\") ]\n",
    "  for i in range(electra_config.num_hidden_layers):\n",
    "    groups.append( list_parameters(model, f\"{base}.encoder.{layers_name}[{i}]\") )\n",
    "  groups.append( list_parameters(model, output_name) )\n",
    "  if electra_config.hidden_size != electra_config.embedding_size:\n",
    "    groups[0] += list_parameters(model, f\"{base}.{scaler_name}\")\n",
    "  if c.my_model and hparam['pre_norm']:\n",
    "    groups[-2] += list_parameters(model, f\"{base}.encoder.norm\")\n",
    "\n",
    "  assert len(list(model.parameters())) == sum([ len(g) for g in groups])\n",
    "  for i, (p1, p2) in enumerate(zip(model.parameters(), [ p for g in groups for p in g])):\n",
    "    assert torch.equal(p1, p2), f\"The {i} th tensor\"\n",
    "  return groups\n",
    "\n",
    "def get_layer_lrs(lr, decay_rate, num_hidden_layers):\n",
    "  lrs = [ lr * (decay_rate ** depth) for depth in range(num_hidden_layers+2)]\n",
    "  if c.original_lr_layer_decays:\n",
    "    for i in range(1, len(lrs)): lrs[i] *= decay_rate\n",
    "  return list(reversed(lrs))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-Mcjq0b8wbkO"
   },
   "source": [
    "## 2.3 learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_glue_learner(task, run_name=None, inference=False):\n",
    "  is_wsc_trick = task=='wnli' and c.wsc_trick\n",
    "\n",
    "  # Num_epochs\n",
    "  if task in ['rte', 'stsb']: num_epochs = 10\n",
    "  else: num_epochs = 3\n",
    "  \n",
    "  # Dataloaders\n",
    "  dls = glue_dls[task]\n",
    "  if isinstance(c.device, str): dls.to(torch.device(c.device))\n",
    "  elif isinstance(c.device, list): dls.to(torch.device('cuda', c.device[0]))\n",
    "  else: dls.to(torch.device('cuda:0'))\n",
    "\n",
    "  # Load pretrained model\n",
    "  if not c.pretrained_checkpoint:\n",
    "    discriminator = ElectraForPreTraining.from_pretrained(f\"google/electra-{c.size}-discriminator\")\n",
    "  else:\n",
    "    discriminator = ModelForDiscriminator(hparam) if c.my_model else ElectraForPreTraining(electra_config)\n",
    "    load_part_model(c.pretrained_ckp_path, discriminator, 'discriminator')\n",
    "\n",
    "  # Seeds & PyTorch benchmark\n",
    "  torch.backends.cudnn.benchmark = True\n",
    "  if c.seeds:\n",
    "    dls[0].rng = random.Random(c.seeds[i]) # for fastai dataloader\n",
    "    random.seed(c.seeds[i])\n",
    "    np.random.seed(c.seeds[i])\n",
    "    torch.manual_seed(c.seeds[i])\n",
    "\n",
    "  # Create finetuning model\n",
    "  if is_wsc_trick: \n",
    "    model = ELECTRAWSCTrickModel(discriminator, hf_tokenizer.pad_token_id)\n",
    "  else:\n",
    "    model = SentencePredictor(discriminator.electra, electra_config.hidden_size, num_class=NUM_CLASS[task])\n",
    "\n",
    "  # Discriminative learning rates\n",
    "  splitter = partial( hf_electra_param_splitter, wsc_trick=is_wsc_trick)\n",
    "  layer_lrs = get_layer_lrs(lr=c.lr, \n",
    "                            decay_rate=c.layer_lr_decay,\n",
    "                            num_hidden_layers=electra_config.num_hidden_layers,)\n",
    "  \n",
    "  # Optimizer\n",
    "  if c.adam_bias_correction: opt_func = partial(Adam, eps=1e-6, mom=0.9, sqr_mom=0.999, wd=c.weight_decay)\n",
    "  else: opt_func = partial(Adam_no_bias_correction, eps=1e-6, mom=0.9, sqr_mom=0.999, wd=c.weight_decay)\n",
    "  \n",
    "  # Learner\n",
    "  learn = Learner(dls, model,\n",
    "                  loss_func=LOSS_FUNC[task], \n",
    "                  opt_func=opt_func,\n",
    "                  metrics=METRICS[task],\n",
    "                  splitter=splitter if not inference else trainable_params,\n",
    "                  lr=layer_lrs if not inference else defaults.lr,\n",
    "                  path='./checkpoints/glue',\n",
    "                  model_dir=c.group_name,)\n",
    "\n",
    "  # Multi gpu\n",
    "  if isinstance(c.device, list) or c.device is None:\n",
    "    learn.create_opt()\n",
    "    learn.model = nn.DataParallel(learn.model, device_ids=c.device)\n",
    "\n",
    "  # Mixed precision\n",
    "  learn.to_native_fp16(init_scale=2.**14)\n",
    "\n",
    "  # Gradient clip\n",
    "  learn.add_cb(GradientClipping(1.0))\n",
    "  \n",
    "  # Logging\n",
    "  # Logging\n",
    "  if run_name and not inference:\n",
    "    if c.logger == 'neptune':\n",
    "      neptune.create_experiment(name=run_name, params={'task':task, **c, **hparam_update})\n",
    "      learn.add_cb(LightNeptuneCallback(False))\n",
    "    elif c.logger == 'wandb':\n",
    "      wandb_run = wandb.init(name=run_name, project='electra_glue', config={'task':task, **c, **hparam_update}, reinit=True)\n",
    "      learn.add_cb(LightWandbCallback(wandb_run))\n",
    "\n",
    "  # Learning rate schedule\n",
    "  if c.schedule == 'one_cycle': \n",
    "    return learn, partial(learn.fit_one_cycle, n_epoch=num_epochs, lr_max=layer_lrs)\n",
    "  elif c.schedule == 'adjusted_one_cycle':\n",
    "    return learn, partial(learn.fit_one_cycle, n_epoch=num_epochs, lr_max=layer_lrs, div=1e5, pct_start=0.1)\n",
    "  else:\n",
    "    lr_shed_func = linear_warmup_and_then_decay if c.schedule=='separate_linear' else linear_warmup_and_decay\n",
    "    lr_shedule = ParamScheduler({'lr': partial(lr_shed_func,\n",
    "                                               lr_max=np.array(layer_lrs),\n",
    "                                               warmup_pct=0.1,\n",
    "                                               total_steps=num_epochs*(len(dls.train)))})\n",
    "    return learn, partial(learn.fit, n_epoch=num_epochs, cbs=[lr_shedule])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 Do finetuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if c.do_finetune:\n",
    "  for i in range(c.start, c.end):\n",
    "    for task in ['cola', 'sst2', 'mrpc', 'stsb', 'rte', 'qnli', 'qqp', 'mnli', 'wnli']:\n",
    "      if c.group_name: run_name = f\"{c.group_name}_{task}_{i}\";\n",
    "      else: run_name = None; print(task)\n",
    "      learn, fit_fc = get_glue_learner(task, run_name)\n",
    "      fit_fc()\n",
    "      if run_name: learn.save(f\"{task}_{i}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "c73X4fdrmnlP"
   },
   "source": [
    "# 3. Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Haven't found way to validate and log two datasets in the training loop, so validate mnli-mm here as a workaround\n",
    "if not c.do_finetune:\n",
    "  learn, _ = get_glue_learner('mnli', inference=True)\n",
    "  learn.load(f\"mnli_{c.th_run['mnli']}\")\n",
    "  with learn.no_mbar():\n",
    "    print(learn.validate(ds_idx=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_identifier(task, split):\n",
    "  \"Turn task name to official task identifier defined.\"\n",
    "  map = {'cola': 'CoLA', 'sst2':'SST-2', 'mrpc':'MRPC', 'qqp':'QQP', 'stsb':'STS-B', 'qnli':'QNLI', 'rte':'RTE', 'wnli':'WNLI', 'ax':'AX'}\n",
    "  if task =='mnli' and split == 'test_matched': return 'MNLI-m'\n",
    "  elif task == 'mnli' and split == 'test_mismatched': return 'MNLI-mm'\n",
    "  else: return map[task]\n",
    "\n",
    "def predict_test(task, checkpoint, dl_idx):\n",
    "  output_dir = Path(f'./test_outputs/glue/{c.group_name}')\n",
    "  output_dir.mkdir(exist_ok=True)\n",
    "  device = torch.device(c.device)\n",
    "\n",
    "  # Load checkpoint and get predictions\n",
    "  learn, _ = get_glue_learner(task, inference=True)\n",
    "  if task == 'wnli' and c.wsc_trick:\n",
    "    load_model_(learn, checkpoint, merge_out_fc=wsc_trick_merge)\n",
    "  else:\n",
    "    load_model_(learn, checkpoint)\n",
    "  results = learn.get_preds(dl=learn.dls[dl_idx], with_decoded=True)\n",
    "  preds = results[-1] # preds -> (predictions logits, targets, decoded prediction)\n",
    "\n",
    "  # Decode target class index to its class name \n",
    "  if task in ['mnli','ax']:\n",
    "    preds = [ ['entailment','neutral','contradiction'][p] for p in preds]\n",
    "  elif task in ['qnli','rte']: \n",
    "    preds = [ ['entailment','not_entailment'][p] for p in preds ]\n",
    "  elif task == 'wnli' and c.wsc_trick:\n",
    "    preds = preds.to(dtype=torch.long).tolist()\n",
    "  else: preds = preds.tolist()\n",
    "    \n",
    "  # Form test dataframe and save\n",
    "  test_df = pd.DataFrame( {'index':range(len(list(glue_dsets[task].values())[dl_idx])), 'prediction': preds} )\n",
    "  split = list(glue_dsets['mnli'].keys())[dl_idx] if task == 'mnli' else 'test'\n",
    "  identifier = get_identifier(task, split)\n",
    "  test_df.to_csv( output_dir/f'{identifier}.tsv', sep='\\t' )\n",
    "  return test_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "wnli\nSome weights of the model checkpoint at google/electra-large-discriminator were not used when initializing ElectraForPreTraining: [&#39;electra.embeddings_project.weight&#39;, &#39;electra.embeddings_project.bias&#39;]\n- This IS expected if you are initializing ElectraForPreTraining from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n- This IS NOT expected if you are initializing ElectraForPreTraining from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": ""
     },
     "metadata": {}
    }
   ],
   "source": [
    "if not c.do_finetune:\n",
    "  for task, th in c.th_run.items():\n",
    "    print(task)\n",
    "    # ax use mnli ckp\n",
    "    if isinstance(th, int):\n",
    "      ckp = f\"{task}_{th}\" if task != 'ax' else f\"mnli_{th}\"\n",
    "    else:\n",
    "      ckp = [f\"{task}_{i}\" if task != 'ax' else f\"mnli_{i}\" for i in th]\n",
    "    # run test for all testset in this task\n",
    "    dl_idxs = [-1, -2] if task=='mnli' else [-1]\n",
    "    for dl_idx in dl_idxs:\n",
    "      df = predict_test(task, ckp, dl_idx)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "yKxetBkqfdgm",
    "MqtDnxK9ZYQP",
    "LFcM5MGaafFJ",
    "7cR6NmFsxH10",
    "n1vFihXdx7iQ",
    "pZS8owurxXUV",
    "admPAZ2M2uqO",
    "1NSoV-Np6utL"
   ],
   "name": "Finetune GLUE with fastai.ipynb",
   "provenance": []
  },
  "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.7-final"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "007556af2af8413ca47aa5f7f725e68a": {
     "model_module": "@jupyter-widgets/controls",
     "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": "initial"
     }
    },
    "2d26c89a3d97493e911c52fa20cf79d9": {
     "model_module": "@jupyter-widgets/controls",
     "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_bd8da3c0b2784f2094d3b0cfa744cf87",
      "placeholder": "​",
      "style": "IPY_MODEL_5b8f1e5172204ed58bd60972289b24a8",
      "value": " 29.0k/29.0k [00:00&lt;00:00, 39.1kB/s]"
     }
    },
    "5b8f1e5172204ed58bd60972289b24a8": {
     "model_module": "@jupyter-widgets/controls",
     "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": ""
     }
    },
    "65eee3543a9d46adbc8037995c3a69a0": {
     "model_module": "@jupyter-widgets/controls",
     "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": "initial"
     }
    },
    "71beed6de71c4aa2862ef33e21cc58bc": {
     "model_module": "@jupyter-widgets/controls",
     "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": ""
     }
    },
    "8d4771f53c7046f1afe7e04a3a647a31": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "9fc29a98d6be46e0974d1f0f7681d5c3": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "a07124360f4c4249a745891d83e719f6": {
     "model_module": "@jupyter-widgets/controls",
     "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": "Downloading: 100%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_8d4771f53c7046f1afe7e04a3a647a31",
      "max": 30329,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_007556af2af8413ca47aa5f7f725e68a",
      "value": 30329
     }
    },
    "a8b32129b9ab41d0895cb1ec445db9ca": {
     "model_module": "@jupyter-widgets/controls",
     "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": "Downloading: 100%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_9fc29a98d6be46e0974d1f0f7681d5c3",
      "max": 28998,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_65eee3543a9d46adbc8037995c3a69a0",
      "value": 28998
     }
    },
    "ae0863b074ff4a1a959514881fe2802f": {
     "model_module": "@jupyter-widgets/controls",
     "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_a8b32129b9ab41d0895cb1ec445db9ca",
       "IPY_MODEL_2d26c89a3d97493e911c52fa20cf79d9"
      ],
      "layout": "IPY_MODEL_e5d5d86746ec4e2298028c57ef775d84"
     }
    },
    "bd8da3c0b2784f2094d3b0cfa744cf87": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "e1d8bdd9c1a74bc18a103bc900589c2c": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "e5d5d86746ec4e2298028c57ef775d84": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "f26a4f6d9334465d9022ab2a3401035b": {
     "model_module": "@jupyter-widgets/controls",
     "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_e1d8bdd9c1a74bc18a103bc900589c2c",
      "placeholder": "​",
      "style": "IPY_MODEL_71beed6de71c4aa2862ef33e21cc58bc",
      "value": " 30.3k/30.3k [00:00&lt;00:00, 392kB/s]"
     }
    },
    "f35afd9d13c240aba4e0619cca0199b7": {
     "model_module": "@jupyter-widgets/base",
     "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
     }
    },
    "f74af04bf4764b7888f0a45e37ce85c1": {
     "model_module": "@jupyter-widgets/controls",
     "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_a07124360f4c4249a745891d83e719f6",
       "IPY_MODEL_f26a4f6d9334465d9022ab2a3401035b"
      ],
      "layout": "IPY_MODEL_f35afd9d13c240aba4e0619cca0199b7"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}