RoBERTa_Fine_Tuning_Emotion_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wj6eoKzotv5I"
   },
   "source": [
    "## Emotion Classification using Fine-tuned BERT model\n",
    "\n",
    "In this tutorial, I will show to fine-tune a language model (LM) for emotion classification with code adapted from this [tutorial](https://zablo.net/blog/post/custom-classifier-on-bert-model-guide-polemo2-sentiment-analysis/) by MARCIN ZABŁOCKI. I adapted his tutorial and modified the code to suit the emotion classification task using a different BERT model. Please refer to his tutorial for more detailed explanations for each code block. I really liked his tutorial because of the attention to detail and the use of high-level libraries to take care of certain parts of the model such as training and finding a good learning rate.\n",
    "\n",
    "Before you get started, make sure to enable `GPU` in the runtime and be sure to\n",
    "restart the runtime in this environment after installing the `pytorch-lr-finder` library.\n",
    "\n",
    "This tutorial is in a rough draft so if you find any issues with this tutorial or have any further questions reach out to me via [Twitter](https://twitter.com/omarsar0).\n",
    "\n",
    "Note that the notebook was created a little while back so if something break it's because the code is not compatible with the library changes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "G2tokZqttmTA",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %%capture\n",
    "# !pip install transformers tokenizers pytorch-lightning torch-lr-finder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I0jZnNegGhZj"
   },
   "source": [
    "Note: you need to Restart runtime after running this code segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "k9ZKIIGvuW5m",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %%capture\n",
    "# !git clone https://github.com/davidtvs/pytorch-lr-finder.git && cd pytorch-lr-finder && python setup.py install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "qqRRWe4UuuIh",
    "outputId": "8cae09a2-9312-441b-ef22-9cd650ae2a1e",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from typing import List\n",
    "import torch.nn.functional as F\n",
    "from transformers import DistilBertTokenizer, AutoTokenizer, AutoModelWithLMHead, DistilBertForSequenceClassification, AdamW, get_linear_schedule_with_warmup\n",
    "import logging\n",
    "import os\n",
    "from functools import lru_cache\n",
    "from tokenizers import ByteLevelBPETokenizer\n",
    "from tokenizers.processors import BertProcessing\n",
    "import pytorch_lightning as pl\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import pandas as pd\n",
    "from argparse import Namespace\n",
    "from sklearn.metrics import classification_report\n",
    "torch.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_whSBDujRiga"
   },
   "source": [
    "## Load the Pretrained Language Model\n",
    "We are first going to look at pretrained language model provided by HuggingFace models. We will use a variant of BERT, called DistilRoBERTa base. The `base` model has less parameters than the `larger` model.\n",
    "\n",
    "[RoBERTa](https://arxiv.org/abs/1907.11692) is a variant of of BERT which \"*modifies key hyperparameters, removing the next-sentence pretraining objective and training with much larger mini-batches and learning rates*\".\n",
    "\n",
    "Knowledge distillation help to train smaller LMs with similar performance and potential."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BvHNcMckSR4M"
   },
   "source": [
    "First, let's load the tokenizer for this model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BPbTd5lmuzQn",
    "outputId": "27a98d20-e967-4dd5-c85d-e69fa78bd091",
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('distilroberta-base')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7KAbKMqJSWRo"
   },
   "source": [
    "Now let's load the actual model with the LM head that takes care of the prediciton for the LM. When fine-tuning we don't use the head and instead use the base model. The code below shows how to do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PCXYlMydzQlP",
    "outputId": "fa0fd766-4a67-46f4-a235-8d943ccfff88",
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = AutoModelWithLMHead.from_pretrained(\"distilroberta-base\")\n",
    "base_model = model.base_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K2_8S8BXSpNa"
   },
   "source": [
    "Let's now try out the tokenizer first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5fidSmH-zrY_",
    "outputId": "fd9d426e-efe2-4679-9949-524e85076c19",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['input_ids', 'attention_mask'])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"Elvis is the king of rock!\"\n",
    "enc = tokenizer.encode_plus(text)\n",
    "enc.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "m8F8yQCDTDQi",
    "outputId": "a6157a68-5e48-460a-a9c5-9ebe1250e798",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': [0, 9682, 9578, 16, 5, 8453, 9, 3152, 328, 2], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n"
     ]
    }
   ],
   "source": [
    "print(enc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P3wSCLKW0ndh"
   },
   "source": [
    "`input_ids` are the numerical encoding of the tokens in the vocabulary. `attention_mask` is an addition option used when batching sequences together and you want to tell the model which tokens should be attented to ([read more](https://huggingface.co/transformers/glossary.html#attention-mask)). The attention mask information helps when dealing with variance in the size of sequences and we need a way to tell the model that we don't want to attend to the padded indices of the sequence.\n",
    "\n",
    "We are only using `input_ids` and `attention_mask`\n",
    "\n",
    "We need to also unsqueeze to simulate batch processing\n",
    "\n",
    "Using DistilBertForSequenceClassification: https://huggingface.co/transformers/model_doc/distilbert.html#distilbertforsequenceclassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Mxsts4uT0PgA",
    "outputId": "07a604b2-2aef-45e6-ef53-1460a2779c6b",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 10, 768])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = base_model(torch.tensor(enc[\"input_ids\"]).unsqueeze(0), torch.tensor(enc[\"attention_mask\"]).unsqueeze(0))\n",
    "out[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZiCO-n_1AHIf",
    "outputId": "0c265024-0729-46c2-f27c-29e3ba473d33",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 768])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## size of representation of one of the tokens\n",
    "out[0][:,0,:].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "srwIb9nr4g4t"
   },
   "source": [
    "`torch.Size([1, 768])` represents batch_size, number of tokens in input text (lenght of tokenized text), model's output hidden size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iAsg0H6g53Bf",
    "outputId": "25d86bd5-02ef-4b52-c1f0-523fb5b4896a",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 9682, 9578, 16, 5, 8453, 9, 3152, 2]\n",
      "<s>Elvis is the king of rock</s>\n",
      "Length: 9\n",
      "torch.Size([9, 768])\n"
     ]
    }
   ],
   "source": [
    "t = \"Elvis is the king of rock\"\n",
    "enc = tokenizer.encode_plus(t)\n",
    "token_representations = base_model(torch.tensor(enc[\"input_ids\"]).unsqueeze(0))[0][0]\n",
    "print(enc[\"input_ids\"])\n",
    "print(tokenizer.decode(enc[\"input_ids\"]))\n",
    "print(f\"Length: {len(enc['input_ids'])}\")\n",
    "print(token_representations.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9RFifOoY7Hsc"
   },
   "source": [
    "## Building Custom Classification head on top of LM base model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vSUMm4Oq7nvR"
   },
   "source": [
    "Use Mish activiation function as in the one proposed in the original tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "tCEDXLxq628O",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from https://github.com/digantamisra98/Mish/blob/b5f006660ac0b4c46e2c6958ad0301d7f9c59651/Mish/Torch/mish.py\n",
    "@torch.jit.script\n",
    "def mish(input):\n",
    "    return input * torch.tanh(F.softplus(input))\n",
    "\n",
    "class Mish(nn.Module):\n",
    "    def forward(self, input):\n",
    "        return mish(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C6Ln6KWm74ku"
   },
   "source": [
    "The model we will use to do the fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "9VDRSRsc71H2",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EmoModel(nn.Module):\n",
    "    def __init__(self, base_model, n_classes, base_model_output_size=768, dropout=0.05):\n",
    "        super().__init__()\n",
    "        self.base_model = base_model\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(base_model_output_size, base_model_output_size),\n",
    "            Mish(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(base_model_output_size, n_classes)\n",
    "        )\n",
    "\n",
    "        for layer in self.classifier:\n",
    "            if isinstance(layer, nn.Linear):\n",
    "                layer.weight.data.normal_(mean=0.0, std=0.02)\n",
    "                if layer.bias is not None:\n",
    "                    layer.bias.data.zero_()\n",
    "\n",
    "    def forward(self, input_, *args):\n",
    "        X, attention_mask = input_\n",
    "        hidden_states = self.base_model(X, attention_mask=attention_mask)\n",
    "\n",
    "        # maybe do some pooling / RNNs... go crazy here!\n",
    "\n",
    "        # use the <s> representation\n",
    "        return self.classifier(hidden_states[0][:, 0, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wjgME-3O8Yfo"
   },
   "source": [
    "### Pretest the model with dummy text\n",
    "We want to ensure that the model is returing the right information back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Y6H9eF8A8XeV",
    "outputId": "9539c3cd-5e2f-4acc-b36a-ab3bb4617dcb",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilroberta-base were not used when initializing RobertaForMaskedLM: ['roberta.pooler.dense.bias', 'roberta.pooler.dense.weight']\n",
      "- This IS expected if you are initializing RobertaForMaskedLM 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 RobertaForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "classifier = EmoModel(AutoModelWithLMHead.from_pretrained(\"distilroberta-base\").base_model, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "-sjfHJ_L9iNH",
    "tags": []
   },
   "outputs": [],
   "source": [
    "X = torch.tensor(enc[\"input_ids\"]).unsqueeze(0).to('cpu')\n",
    "attn = torch.tensor(enc[\"attention_mask\"]).unsqueeze(0).to('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "o6QhCuEC-y2z",
    "outputId": "ee9d2df4-9f93-4cf1-a21d-71613650e3ab",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1118, -0.1068,  0.0190]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier((X, attn))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I-N7WSY7Cb7v"
   },
   "source": [
    "## Prepare your dataset for fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jDWkjaLV-5tj",
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -p tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wMMm5Ye1Db-m",
    "outputId": "9de6db3c-814f-4f31-e669-a6f9f574b27b",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('tokenizer/tokenizer_config.json',\n",
       " 'tokenizer/special_tokens_map.json',\n",
       " 'tokenizer/vocab.json',\n",
       " 'tokenizer/merges.txt',\n",
       " 'tokenizer/added_tokens.json',\n",
       " 'tokenizer/tokenizer.json')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## load pretrained tokenizer information\n",
    "tokenizer.save_pretrained(\"tokenizer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3FVtbmrzDkF8",
    "outputId": "ce825928-4005-4967-f866-b98471ce16b8",
    "tags": []
   },
   "outputs": [],
   "source": [
    "!ls tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BhTEgIaLEDRo"
   },
   "source": [
    "Implement CollateFN using fast tokenizers.\n",
    "This function basically takes care of proper tokenization and batches of sequences. This way you don't need to create your batches manually. Find out more about Tokenizers [here](https://github.com/huggingface/tokenizers/tree/master/bindings/python)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3SCLBZsMDn4s",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TokenizersCollateFn:\n",
    "    def __init__(self, max_tokens=512):\n",
    "\n",
    "        ## RoBERTa uses BPE tokenizer similar to GPT\n",
    "        t = ByteLevelBPETokenizer(\n",
    "            \"tokenizer/vocab.json\",\n",
    "            \"tokenizer/merges.txt\"\n",
    "        )\n",
    "        t._tokenizer.post_processor = BertProcessing(\n",
    "            (\"</s>\", t.token_to_id(\"</s>\")),\n",
    "            (\"<s>\", t.token_to_id(\"<s>\")),\n",
    "        )\n",
    "        t.enable_truncation(max_tokens)\n",
    "        t.enable_padding(length=max_tokens, pad_id=t.token_to_id(\"<pad>\"))\n",
    "        self.tokenizer = t\n",
    "\n",
    "    def __call__(self, batch):\n",
    "        encoded = self.tokenizer.encode_batch([x['text'] for x in batch])\n",
    "        sequences_padded = torch.tensor([enc.ids for enc in encoded])\n",
    "        attention_masks_padded = torch.tensor([enc.attention_mask for enc in encoded])\n",
    "        labels = torch.tensor([x['label'] for x in batch])\n",
    "\n",
    "        return (sequences_padded, attention_masks_padded), labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4hu70Ng0Eqls"
   },
   "source": [
    "## Getting the Data and Preview it\n",
    "Below we are going to load the data and show you how to create the splits. However, we don't need to split the data manually becuase I have already created the splits and stored those files seperately which you can quickly download below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds = load_dataset('emotion')\n",
    "ds['validation'][19]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "r_03fxufWX_G",
    "tags": []
   },
   "outputs": [],
   "source": [
    "## emotion labels\n",
    "label2int = {\n",
    "  \"sadness\": 0,\n",
    "  \"joy\": 1,\n",
    "  \"love\": 2,\n",
    "  \"anger\": 3,\n",
    "  \"fear\": 4,\n",
    "  \"surprise\": 5\n",
    "}\n",
    "\n",
    "emotions = [ \"sadness\", \"joy\", \"love\", \"anger\", \"fear\", \"surprise\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rAD1J6c0dLp8"
   },
   "source": [
    "## Create the Dataset object"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aOOI69vwIYcN"
   },
   "source": [
    "Create the Dataset object that will be used to load the different datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "Ktr6xeMuISin",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EmoDataset(Dataset):\n",
    "    def __init__(self, path):\n",
    "        super().__init__()\n",
    "        self.data_column = \"text\"\n",
    "        self.class_column = \"class\"\n",
    "        self.data = pd.read_csv(path, sep=\";\", header=None, names=[self.data_column, self.class_column],\n",
    "                               engine=\"python\")\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data.loc[idx, self.data_column], label2int[self.data.loc[idx, self.class_column]]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0h6tTn9hd6v8"
   },
   "source": [
    "## Training with PyTorchLightning\n",
    "\n",
    "[PyTorchLightning](https://www.pytorchlightning.ai/) is a library that abstracts the complexity of training neural networks with PyTorch. It is built on top of PyTorch and simplifies training.\n",
    "\n",
    "![](https://pytorch-lightning.readthedocs.io/en/latest/_images/pt_to_pl.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "zlMbsLMdOufN",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch_lr_finder import LRFinder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "LlDd6a-INsBR",
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_path = \"train\"\n",
    "test_path = \"test\"\n",
    "val_path = \"validation\"\n",
    "val_path = \"test\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "# from codes import DictEnum, auto\n",
    "\n",
    "\n",
    "# class Optimizer(DictEnum):\n",
    "#     ADAM = auto()\n",
    "#     KATE = auto()\n",
    "\n",
    "\n",
    "class ADAM(torch.optim.Optimizer):\n",
    "    def __init__(self, params, lr):  # lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):\n",
    "        defaults = dict(lr=lr)\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "\n",
    "    def step(self):\n",
    "        loss = None\n",
    "\n",
    "        for group in self.param_groups:\n",
    "\n",
    "            for p in group['params']:\n",
    "                grad = p.grad.data\n",
    "                state = self.state[p]\n",
    "\n",
    "                # State initialization\n",
    "                if len(state) == 0:\n",
    "                    state['step'] = 0\n",
    "                    # Momentum (Exponential MA of gradients)\n",
    "                    state['m'] = torch.zeros_like(p.data)\n",
    "\n",
    "                    # RMS Prop componenet. (Exponential MA of squared gradients). Denominator.\n",
    "                    state['v'] = torch.zeros_like(p.data)\n",
    "\n",
    "                m, v = state['m'], state['v']\n",
    "\n",
    "                beta1, beta2 = 0.9, 0.999\n",
    "                state['step'] += 1\n",
    "\n",
    "                # Add weight decay if any\n",
    "                # if group['weight_decay'] != 0:\n",
    "                #     grad = grad + group['weight_decay']*p.data\n",
    "\n",
    "                # Momentum\n",
    "                m = torch.mul(m, beta1) + (1 - beta1)*grad\n",
    "\n",
    "                # RMS\n",
    "                v = torch.mul(v, beta2) + (1-beta2)*(grad*grad)\n",
    "\n",
    "                mhat = m / (1 - beta1 ** state['step'])\n",
    "                vhat = v / (1 - beta2 ** state['step'])\n",
    "\n",
    "                eps = 1e-4\n",
    "                denom = torch.sqrt(vhat) + eps\n",
    "\n",
    "                lr = group['lr']\n",
    "                p.data = p.data - lr * mhat / denom\n",
    "\n",
    "                # Save state\n",
    "                state['m'], state['v'] = m, v\n",
    "\n",
    "        return loss\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class KATE(torch.optim.Optimizer): # delta 0 or 1e-8\n",
    "    def __init__(self, params, lr):  # lr=1e-3, eta=0.9, eps=1e-8, delta=0, weight_decay=0):\n",
    "        defaults = dict(lr=lr)\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "    def step(self):\n",
    "        loss = None\n",
    "        for group in self.param_groups:\n",
    "\n",
    "            for p in group['params']:\n",
    "                grad = p.grad.data\n",
    "                state = self.state[p]\n",
    "\n",
    "                # State initialization\n",
    "                if len(state) == 0:\n",
    "                    state['step'] = 0\n",
    "                    state['m'] = torch.zeros_like(p.data)\n",
    "                    state['b'] = torch.zeros_like(p.data)\n",
    "\n",
    "                m, b = state['m'], state['b']\n",
    "                eta = 0.001\n",
    "\n",
    "                # if group['weight_decay'] != 0:\n",
    "                #     grad = grad + group['weight_decay']*p.data\n",
    "\n",
    "\n",
    "                g = grad*grad\n",
    "\n",
    "                b = b + g\n",
    "                eps = 1e-5\n",
    "                denom = b + eps\n",
    "                m = m + torch.mul(eta, g) + g / denom\n",
    "\n",
    "                lr = group['lr']\n",
    "                p.data = p.data - lr * torch.sqrt(m) * grad / denom\n",
    "\n",
    "                # Save state\n",
    "                state['m'], state['b'] = m, b\n",
    "                state['step'] += 1\n",
    "\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from mlflow import MlflowClient\n",
    "from mlflow.entities import ViewType\n",
    "\n",
    "\n",
    "class MLFlowLogger():\n",
    "    def __init__(self):\n",
    "        tracking_uri = os.path.expanduser('~/mlruns/')\n",
    "        experiment_name = os.environ['MLFLOW_EXPERIMENT_NAME']\n",
    "        self.c = MlflowClient(tracking_uri=tracking_uri)\n",
    "        self.e = self.c.get_experiment_by_name(experiment_name)\n",
    "        self.e_id = self.c.create_experiment(experiment_name)\\\n",
    "            if self.e is None else self.e.experiment_id\n",
    "\n",
    "        self.enabled = True\n",
    "\n",
    "    def check_exist(self, config):\n",
    "        check_exist = eval(os.environ['MLFLOW_CHECK_EXIST'])\n",
    "        if not check_exist:\n",
    "            return False\n",
    "\n",
    "        if self.e is not None:\n",
    "            filter_string = list()\n",
    "            for key in config.keys():\n",
    "                if isinstance(config[key], dict):\n",
    "                    value = config[key]['name']\n",
    "                else:\n",
    "                    value = config[key]\n",
    "                filter_string.append(f'params.{key}=\"{value}\"')\n",
    "            filter_string.append('attributes.status=\"FINISHED\"')\n",
    "            tags = eval(os.environ['MLFLOW_RUN_TAGS'])\n",
    "            for key, value in tags.items():\n",
    "                filter_string.append(f'tags.\"{key}\" = \"{value}\"')\n",
    "            # print(f\"{filter_string=}\")\n",
    "            filter_string = ' and '.join(filter_string)\n",
    "            runs = self.c.search_runs(experiment_ids=[self.e.experiment_id],\n",
    "                                      filter_string=filter_string,\n",
    "                                      run_view_type=ViewType.ACTIVE_ONLY)\n",
    "            if len(runs):\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    def init(self, config):\n",
    "        if not self.enabled:\n",
    "            return\n",
    "\n",
    "        r = self.c.create_run(experiment_id=self.e_id,\n",
    "                              run_name=os.environ['MLFLOW_RUN_NAME'],\n",
    "                              tags=eval(os.environ['MLFLOW_RUN_TAGS']))\n",
    "        self.r_id = r.info.run_id\n",
    "        self.c.log_dict(self.r_id, config, 'config.json')\n",
    "        for key, value in config.items():\n",
    "            if isinstance(value, dict):\n",
    "                value = value['name']\n",
    "            self.c.log_param(self.r_id, key, value)\n",
    "\n",
    "    def log_metrics(self, metrics, step):\n",
    "        if not self.enabled:\n",
    "            return\n",
    "\n",
    "        for key in metrics.keys():\n",
    "            self.c.log_metric(self.r_id, key, metrics[key], step=step)\n",
    "        # print('Step ' + '{}: '.format(step).rjust(6) +\n",
    "              # ' '.join(\"{}: {:.5f}\".format(k, v) for k, v in metrics.items() if 'weights' not in k), end='\\r', flush=True)\n",
    "\n",
    "    def terminate(self):\n",
    "        if not self.enabled:\n",
    "            return\n",
    "\n",
    "        self.c.set_terminated(self.r_id)\n",
    "        print()\n",
    "\n",
    "os.environ['MLFLOW_VERBOSE'] = 'True'\n",
    "# os.environ['MLFLOW_CHECK_EXIST'] = 'False'\n",
    "os.environ['MLFLOW_CHECK_EXIST'] = 'True'\n",
    "os.environ['MLFLOW_EXPERIMENT_NAME'] = os.path.basename(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "RJHhNRcZK7sV",
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Methods required by PyTorchLightning\n",
    "\n",
    "class TrainingModule(pl.LightningModule):\n",
    "    def __init__(self, hparams):\n",
    "        super().__init__()\n",
    "        self.model = EmoModel(AutoModelWithLMHead.from_pretrained(\"distilroberta-base\").base_model, len(emotions))\n",
    "        self.loss = nn.CrossEntropyLoss() ## combines LogSoftmax() and NLLLoss()\n",
    "        #self.hparams = hparams\n",
    "        self.hparams.update(vars(hparams))\n",
    "        self.test_step_y_hats = []\n",
    "        self.test_step_ys = []\n",
    "\n",
    "        config = vars(hparams)\n",
    "        self.loggger = MLFlowLogger()\n",
    "        if self.loggger.check_exist(config):\n",
    "            return\n",
    "        self.loggger.enabled = eval(os.environ['MLFLOW_VERBOSE'])\n",
    "        self.loggger.init(config)\n",
    "        self.loggger.step = 0\n",
    "\n",
    "\n",
    "    def step(self, batch, step_name=\"train\"):\n",
    "        X, y = batch\n",
    "        pred = self.forward(X)\n",
    "        loss = self.loss(pred, y)\n",
    "        # loss_key = f\"{step_name}_loss\"\n",
    "        # tensorboard_logs = {loss_key: loss}\n",
    "\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "        # return { (\"loss\" if step_name == \"train\" else loss_key): loss, 'log': tensorboard_logs,\n",
    "               # \"progress_bar\": {loss_key: loss}}\n",
    "        # return pred, y\n",
    "\n",
    "    def forward(self, X, *args):\n",
    "        return self.model(X, *args)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        return self.step(batch, \"train\")\n",
    "\n",
    "    # def validation_step(self, batch, batch_idx):\n",
    "    #     # return self.step(batch, \"val\")\n",
    "    #     X, y = batch\n",
    "    #     pred = self.forward(X)\n",
    "    #     loss = self.loss(pred, y)\n",
    "        # self.log(\"val_loss\", val_loss)\n",
    "\n",
    "    # def validation_end(self, outputs: List[dict]):\n",
    "    #     loss = torch.stack([x[\"val_loss\"] for x in outputs]).mean()\n",
    "    #     # return {\"val_loss\": loss}\n",
    "\n",
    "    # def test_step(self, batch, batch_idx):\n",
    "    #     x, y = batch\n",
    "    #     logits = self(x)\n",
    "    #     # self.test_acc(logits, y)\n",
    "    #     # self.log('test_acc', self.test_acc, on_step=False, on_epoch=True)\n",
    "    #     preds = logits.argmax(dim=-1)\n",
    "    #     acc = (y == preds).float().mean()\n",
    "    #     return acc\n",
    "    #     # print(acc)\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        return self.step(batch, \"test\")\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        X, y = batch\n",
    "        logits = self.forward(X)\n",
    "        loss = self.loss(logits, y)\n",
    "#         loss_key = f\"{step_name}_loss\"\n",
    "#         tensorboard_logs = {loss_key: loss}\n",
    "\n",
    "#         return { (\"loss\" if step_name == \"train\" else loss_key): loss, 'log': tensorboard_logs,\n",
    "#                \"progress_bar\": {loss_key: loss}}\n",
    "        # return pred, y\n",
    "        self.log(\"val_loss\", loss)\n",
    "        self.test_step_y_hats.append(logits)\n",
    "        self.test_step_ys.append(y)\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        # all_outputs is a list of [targets, logits] tuples from each test.\n",
    "        acc = 0\n",
    "        for logits, y in zip(self.test_step_y_hats, self.test_step_ys):\n",
    "            preds = logits.argmax(dim=-1)\n",
    "            acc += (y == preds).float().mean()\n",
    "            # print(acc)\n",
    "            # do something\n",
    "        # calculate metrics based on all outputs from test.\n",
    "        # log the results.\n",
    "        acc /= len(self.test_step_y_hats)\n",
    "        # print('acc', acc)\n",
    "        # log = {}\n",
    "        # log.update({'Accuracy': acc})\n",
    "        # self.log_dict(log)\n",
    "        self.test_step_y_hats = []\n",
    "        self.test_step_ys = []\n",
    "\n",
    "        if self.loggger.enabled:\n",
    "            self.loggger.log_metrics({'accuracy': acc}, self.loggger.step)\n",
    "            self.loggger.step += 1\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.create_data_loader(self.hparams.train_path, shuffle=True)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.create_data_loader(self.hparams.val_path)\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return self.create_data_loader(self.hparams.test_path)\n",
    "\n",
    "    def create_data_loader(self, ds_path: str, shuffle=False):\n",
    "        return DataLoader(\n",
    "                    ds[ds_path],\n",
    "                    batch_size=self.hparams.batch_size,\n",
    "                    shuffle=shuffle,\n",
    "                    collate_fn=TokenizersCollateFn()\n",
    "        )\n",
    "\n",
    "    @lru_cache()\n",
    "    def total_steps(self):\n",
    "        return len(self.train_dataloader()) // self.hparams.accumulate_grad_batches * self.hparams.epochs\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        ## use AdamW optimizer -- faster approach to training NNs\n",
    "        ## read: https://www.fast.ai/2018/07/02/adam-weight-decay/\n",
    "        # optimizer = AdamW(self.model.parameters(), lr=self.hparams.lr)\n",
    "        # optimizer = KATE(self.model.parameters(), lr=self.hparams.lr)\n",
    "        optimizer = ADAM(self.model.parameters(), lr=self.hparams.lr)\n",
    "        lr_scheduler = get_linear_schedule_with_warmup(\n",
    "                    optimizer,\n",
    "                    num_warmup_steps=self.hparams.warmup_steps,\n",
    "                    num_training_steps=self.total_steps(),\n",
    "        )\n",
    "        return [optimizer], [{\"scheduler\": lr_scheduler, \"interval\": \"step\"}]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OGc7Vw1moHxr"
   },
   "source": [
    "## Finding Learning rate for the model\n",
    "\n",
    "The code below aims to obtain valuable information about the optimal learning rate during a pretraining run. Determine boundary and increase the leanring rate linearly or exponentially.\n",
    "\n",
    "More: https://github.com/davidtvs/pytorch-lr-finder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 628,
     "referenced_widgets": [
      "2474a0f93d7949ef9b28e776d8912bf3",
      "3e5ba44c3fc44d7e9c4efeb514241188",
      "1b9bf3f166364598b0ef8aeaf5920b79",
      "95b619db99ac465088112f459b2b7d9a",
      "981018fbaa6d4087af65abe6b6571b15",
      "d59aafa28cdb4710ad96b4eb61e7d1db",
      "479e8d923467477c9788baff771d4e49",
      "cb6af25cd4f5443d8ee337f47bd45172",
      "fa0ddec00b8b4e59b2a62aafbf99a167",
      "e4bb089a7a4f45bc944c15b753739715",
      "cc0ed6a14a81474da927788fc54d226d"
     ]
    },
    "id": "xL4lNPDFoFyU",
    "outputId": "e4b9592f-0ea3-4d3d-bd87-c99c5b8b0bf4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# lr=0.1 ## uper bound LR\n",
    "# from torch_lr_finder import LRFinder\n",
    "# hparams_tmp = Namespace(\n",
    "#     train_path=train_path,\n",
    "#     val_path=val_path,\n",
    "#     test_path=test_path,\n",
    "#     batch_size=16,\n",
    "#     warmup_steps=100,\n",
    "#     epochs=1,\n",
    "#     lr=lr,\n",
    "#     accumulate_grad_batches=1,\n",
    "# )\n",
    "# module = TrainingModule(hparams_tmp)\n",
    "# criterion = nn.CrossEntropyLoss()\n",
    "# optimizer = AdamW(module.parameters(), lr=5e-7) ## lower bound LR\n",
    "# # lr_finder = LRFinder(module, optimizer, criterion, device=\"cuda\")\n",
    "# # lr_finder = LRFinder(module, optimizer, criterion, device=\"cuda\")\n",
    "# lr_finder.range_test(module.train_dataloader(), end_lr=100, num_iter=100, accumulation_steps=hparams_tmp.accumulate_grad_batches)\n",
    "# lr_finder.plot()\n",
    "# lr_finder.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YdqP56M1oXav",
    "outputId": "e028fba7-57a0-4625-f00d-e56ba294fa74",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0001"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr = 1e-4 #  the value from the original code user for Adam\n",
    "lr\n",
    "\n",
    "# lr = 1e-5 #  stepsize for KATE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 506
    },
    "id": "vMab6vu0Bow0",
    "outputId": "08307b02-3478-40b5-fa5b-bfa8c91ec43d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# lr_finder.plot(show_lr=lr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZhHutCseBxjJ"
   },
   "source": [
    "## Training the Emotion Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q3FiLr3LBrjs",
    "outputId": "4bf02008-b6db-4b08-f623-7392e05494cc",
    "tags": []
   },
   "outputs": [],
   "source": [
    "hparams = Namespace(\n",
    "    train_path=train_path,\n",
    "    val_path=val_path,\n",
    "    test_path=test_path,\n",
    "    batch_size=64,\n",
    "    warmup_steps=100,\n",
    "    epochs=10,\n",
    "    lr=lr,\n",
    "    accumulate_grad_batches=1\n",
    ")\n",
    "os.environ['MLFLOW_RUN_TAGS'] = str(dict(about=f'tuned'))\n",
    "os.environ['MLFLOW_RUN_NAME'] = 'Adam'\n",
    "# os.environ['MLFLOW_RUN_NAME'] = 'KATE'\n",
    "# os.environ['MLFLOW_RUN_NAME'] = 'adam my'\n",
    "module = TrainingModule(hparams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "hEC-bwa4RwzJ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 399,
     "referenced_widgets": [
      "a13b519fa03e4e1786ec380ab2c87524",
      "c1fdb14e408c46b08082d7a8a7b74f4a",
      "f53fe4df386c416d8f4a56174a9e0b61",
      "89ad2e8f1a42442098277737cf8b92aa",
      "6280a6c12b66426f82c28df35397129e",
      "988fc843fb034b0c8c4c7c842ab8c899",
      "7eddf5e592bc48f9ba2e679db37a15c3",
      "5137ad3daf1e43b39dd5c0c3e404852f",
      "8ea8a78a76b449cf9146317fa05f9620",
      "02dfd4fb9dad465c8acc09eaf9a6b285",
      "052c771dae774a74a6c6a8610fc81e34",
      "263975abb6854dac83015c7310833849",
      "fa5beec934d7450ba3df6f90f067e303",
      "09e0798e9774467e98afff55de96d012",
      "4bcb1b9e06594611be5242001af988bf",
      "2bc3d7dd337046baa489c067aa02659a",
      "3286f501c31e4f6e9275df86b57459a5",
      "fc68144e8f6d43c0ab7b2be124708c34",
      "91a34baeefdc471eaf445d0846d4216e",
      "b1f29eec22e14e429b692b9c3c479e89",
      "4b25d7e69c2049c386968d0bbb3d8712",
      "131d891450e245ffb888e28147be5692",
      "bb13587a10224884a479e285d7cd2f0e",
      "c157b85ec1ca48528b1c863ab5d4f7eb",
      "3ddbe060d003491591c3a31540eb9fb9",
      "f55f64c8f18e469e91f43e6f70a45ffa",
      "41da8c57c82a4399afeb0586902564ca",
      "f63d44aa1418469aa51c3c08705a0ca2",
      "26a8f8b9f5c648d59d7e8993dc7bdb53",
      "9d5ba7e666924449b00c5fc92f065fd9",
      "4067bd90f3ae408985e3def50f3377e7",
      "ebb81dbcd075448d90c37cb6365a5220",
      "d65a2074389e402c82dd7364fd1163a0"
     ]
    },
    "id": "oRnl4HXvB5-T",
    "outputId": "00feb06d-bee7-4c39-cac1-354dd677316f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.set_float32_matmul_precision('medium')\n",
    "trainer = pl.Trainer(accelerator=\"auto\", max_epochs=hparams.epochs,\n",
    "                     accumulate_grad_batches=hparams.accumulate_grad_batches, precision='16-mixed', check_val_every_n_epoch=1)\n",
    "trainer.fit(module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "02dfd4fb9dad465c8acc09eaf9a6b285": {
     "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
     }
    },
    "052c771dae774a74a6c6a8610fc81e34": {
     "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": ""
     }
    },
    "09e0798e9774467e98afff55de96d012": {
     "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_91a34baeefdc471eaf445d0846d4216e",
      "max": 500,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_b1f29eec22e14e429b692b9c3c479e89",
      "value": 500
     }
    },
    "131d891450e245ffb888e28147be5692": {
     "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": ""
     }
    },
    "1b9bf3f166364598b0ef8aeaf5920b79": {
     "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": "danger",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_cb6af25cd4f5443d8ee337f47bd45172",
      "max": 100,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_fa0ddec00b8b4e59b2a62aafbf99a167",
      "value": 67
     }
    },
    "2474a0f93d7949ef9b28e776d8912bf3": {
     "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_3e5ba44c3fc44d7e9c4efeb514241188",
       "IPY_MODEL_1b9bf3f166364598b0ef8aeaf5920b79",
       "IPY_MODEL_95b619db99ac465088112f459b2b7d9a"
      ],
      "layout": "IPY_MODEL_981018fbaa6d4087af65abe6b6571b15"
     }
    },
    "263975abb6854dac83015c7310833849": {
     "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_fa5beec934d7450ba3df6f90f067e303",
       "IPY_MODEL_09e0798e9774467e98afff55de96d012",
       "IPY_MODEL_4bcb1b9e06594611be5242001af988bf"
      ],
      "layout": "IPY_MODEL_2bc3d7dd337046baa489c067aa02659a"
     }
    },
    "26a8f8b9f5c648d59d7e8993dc7bdb53": {
     "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": ""
     }
    },
    "2bc3d7dd337046baa489c067aa02659a": {
     "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": "inline-flex",
      "flex": null,
      "flex_flow": "row wrap",
      "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": "100%"
     }
    },
    "3286f501c31e4f6e9275df86b57459a5": {
     "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
     }
    },
    "3ddbe060d003491591c3a31540eb9fb9": {
     "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": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_9d5ba7e666924449b00c5fc92f065fd9",
      "max": 63,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_4067bd90f3ae408985e3def50f3377e7",
      "value": 63
     }
    },
    "3e5ba44c3fc44d7e9c4efeb514241188": {
     "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_d59aafa28cdb4710ad96b4eb61e7d1db",
      "placeholder": "​",
      "style": "IPY_MODEL_479e8d923467477c9788baff771d4e49",
      "value": " 67%"
     }
    },
    "4067bd90f3ae408985e3def50f3377e7": {
     "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": ""
     }
    },
    "41da8c57c82a4399afeb0586902564ca": {
     "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": "inline-flex",
      "flex": null,
      "flex_flow": "row wrap",
      "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": "hidden",
      "width": "100%"
     }
    },
    "479e8d923467477c9788baff771d4e49": {
     "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": ""
     }
    },
    "4b25d7e69c2049c386968d0bbb3d8712": {
     "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
     }
    },
    "4bcb1b9e06594611be5242001af988bf": {
     "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_4b25d7e69c2049c386968d0bbb3d8712",
      "placeholder": "​",
      "style": "IPY_MODEL_131d891450e245ffb888e28147be5692",
      "value": " 500/500 [11:19&lt;00:00,  0.74it/s, v_num=0]"
     }
    },
    "5137ad3daf1e43b39dd5c0c3e404852f": {
     "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": "2",
      "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
     }
    },
    "6280a6c12b66426f82c28df35397129e": {
     "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": "inline-flex",
      "flex": null,
      "flex_flow": "row wrap",
      "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": "hidden",
      "width": "100%"
     }
    },
    "7eddf5e592bc48f9ba2e679db37a15c3": {
     "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": ""
     }
    },
    "89ad2e8f1a42442098277737cf8b92aa": {
     "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_02dfd4fb9dad465c8acc09eaf9a6b285",
      "placeholder": "​",
      "style": "IPY_MODEL_052c771dae774a74a6c6a8610fc81e34",
      "value": " 2/2 [00:01&lt;00:00,  1.97it/s]"
     }
    },
    "8ea8a78a76b449cf9146317fa05f9620": {
     "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": ""
     }
    },
    "91a34baeefdc471eaf445d0846d4216e": {
     "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": "2",
      "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
     }
    },
    "95b619db99ac465088112f459b2b7d9a": {
     "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_e4bb089a7a4f45bc944c15b753739715",
      "placeholder": "​",
      "style": "IPY_MODEL_cc0ed6a14a81474da927788fc54d226d",
      "value": " 67/100 [00:50&lt;00:23,  1.38it/s]"
     }
    },
    "981018fbaa6d4087af65abe6b6571b15": {
     "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
     }
    },
    "988fc843fb034b0c8c4c7c842ab8c899": {
     "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
     }
    },
    "9d5ba7e666924449b00c5fc92f065fd9": {
     "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": "2",
      "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
     }
    },
    "a13b519fa03e4e1786ec380ab2c87524": {
     "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_c1fdb14e408c46b08082d7a8a7b74f4a",
       "IPY_MODEL_f53fe4df386c416d8f4a56174a9e0b61",
       "IPY_MODEL_89ad2e8f1a42442098277737cf8b92aa"
      ],
      "layout": "IPY_MODEL_6280a6c12b66426f82c28df35397129e"
     }
    },
    "b1f29eec22e14e429b692b9c3c479e89": {
     "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": ""
     }
    },
    "bb13587a10224884a479e285d7cd2f0e": {
     "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_c157b85ec1ca48528b1c863ab5d4f7eb",
       "IPY_MODEL_3ddbe060d003491591c3a31540eb9fb9",
       "IPY_MODEL_f55f64c8f18e469e91f43e6f70a45ffa"
      ],
      "layout": "IPY_MODEL_41da8c57c82a4399afeb0586902564ca"
     }
    },
    "c157b85ec1ca48528b1c863ab5d4f7eb": {
     "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_f63d44aa1418469aa51c3c08705a0ca2",
      "placeholder": "​",
      "style": "IPY_MODEL_26a8f8b9f5c648d59d7e8993dc7bdb53",
      "value": "Validation DataLoader 0: 100%"
     }
    },
    "c1fdb14e408c46b08082d7a8a7b74f4a": {
     "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_988fc843fb034b0c8c4c7c842ab8c899",
      "placeholder": "​",
      "style": "IPY_MODEL_7eddf5e592bc48f9ba2e679db37a15c3",
      "value": "Sanity Checking DataLoader 0: 100%"
     }
    },
    "cb6af25cd4f5443d8ee337f47bd45172": {
     "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
     }
    },
    "cc0ed6a14a81474da927788fc54d226d": {
     "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": ""
     }
    },
    "d59aafa28cdb4710ad96b4eb61e7d1db": {
     "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
     }
    },
    "d65a2074389e402c82dd7364fd1163a0": {
     "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": ""
     }
    },
    "e4bb089a7a4f45bc944c15b753739715": {
     "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
     }
    },
    "ebb81dbcd075448d90c37cb6365a5220": {
     "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
     }
    },
    "f53fe4df386c416d8f4a56174a9e0b61": {
     "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": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5137ad3daf1e43b39dd5c0c3e404852f",
      "max": 2,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_8ea8a78a76b449cf9146317fa05f9620",
      "value": 2
     }
    },
    "f55f64c8f18e469e91f43e6f70a45ffa": {
     "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_ebb81dbcd075448d90c37cb6365a5220",
      "placeholder": "​",
      "style": "IPY_MODEL_d65a2074389e402c82dd7364fd1163a0",
      "value": " 63/63 [00:29&lt;00:00,  2.16it/s]"
     }
    },
    "f63d44aa1418469aa51c3c08705a0ca2": {
     "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
     }
    },
    "fa0ddec00b8b4e59b2a62aafbf99a167": {
     "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": ""
     }
    },
    "fa5beec934d7450ba3df6f90f067e303": {
     "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_3286f501c31e4f6e9275df86b57459a5",
      "placeholder": "​",
      "style": "IPY_MODEL_fc68144e8f6d43c0ab7b2be124708c34",
      "value": "Epoch 0: 100%"
     }
    },
    "fc68144e8f6d43c0ab7b2be124708c34": {
     "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": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}