Pytorch_Lightning_BERT_Huggingface_w_SAM_+_EMA.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Pytorch Lightning BERT Huggingface w SAM + EMA.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/maxmatical/ml-cheatsheet/blob/master/Pytorch_Lightning_BERT_Huggingface_w_SAM_%2B_EMA.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nzBFZ3cJG7AV"
      },
      "source": [
        "# Inspirations\n",
        "https://curiousily.com/posts/multi-label-text-classification-with-bert-and-pytorch-lightning/\n",
        "\n",
        "https://github.com/mgrankin/over9000/blob/master/train.py\n",
        "\n",
        "\n",
        "https://medium.com/pytorch/getting-started-with-ray-lightning-easy-multi-node-pytorch-lightning-training-e639031aff8b\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1o9pMmzzoq7e",
        "outputId": "9fc6e270-a880-415a-e04e-e34e1ca5aa49"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sat Jan  1 21:42:21 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 495.44       Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   48C    P0    58W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5QkBmYey_BMu"
      },
      "source": [
        "%%capture\n",
        "!pip install transformers\n",
        "!pip install pytorch_lightning\n",
        "!pip install torchmetrics\n",
        "!pip install torch-ema"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zDygGNnC_VGv"
      },
      "source": [
        "from typing import Optional\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.optim.lr_scheduler import OneCycleLR # , ReduceLROnPlateau\n",
        "from torch.optim import AdamW\n",
        "\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel\n",
        "  \n",
        "import pytorch_lightning as pl\n",
        "from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor, StochasticWeightAveraging\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "import torchmetrics\n",
        "from torchmetrics.functional import accuracy, f1, auroc\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import classification_report, multilabel_confusion_matrix\n",
        "\n",
        "from torch_ema import ExponentialMovingAverage"
      ],
      "execution_count": 42,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SWso8EH2IDcU",
        "outputId": "d708bbdb-89d5-46d7-80cd-d290fe17695d"
      },
      "source": [
        "# pretrained model and tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"distilroberta-base\")\n",
        "model = AutoModel.from_pretrained(\"distilroberta-base\")"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at distilroberta-base were not used when initializing RobertaModel: ['lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.decoder.weight', 'lm_head.bias', 'lm_head.layer_norm.weight', 'lm_head.layer_norm.bias']\n",
            "- This IS expected if you are initializing RobertaModel 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 RobertaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4lqOT-HXGYIN"
      },
      "source": [
        "# Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7WLrlT3NGZGT",
        "outputId": "871cca6e-ae42-4265-f181-7b7adb8002ac"
      },
      "source": [
        "!gdown --id 1VuQ-U7TtggShMeuRSA_hzC8qGDl2LRkr"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1VuQ-U7TtggShMeuRSA_hzC8qGDl2LRkr\n",
            "To: /content/toxic_comments.csv\n",
            "100% 68.8M/68.8M [00:00<00:00, 132MB/s] \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "id": "xuywc9Y5Gbf1",
        "outputId": "9e7a61ab-3997-4b71-d7cb-f863a819890e"
      },
      "source": [
        "df = pd.read_csv(\"toxic_comments.csv\")\n",
        "\n",
        "df.head()"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\n",
              "  <div id=\"df-3f2b0022-91d6-465b-b0c9-128092b45884\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>id</th>\n",
              "      <th>comment_text</th>\n",
              "      <th>toxic</th>\n",
              "      <th>severe_toxic</th>\n",
              "      <th>obscene</th>\n",
              "      <th>threat</th>\n",
              "      <th>insult</th>\n",
              "      <th>identity_hate</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>0000997932d777bf</td>\n",
              "      <td>Explanation\\nWhy the edits made under my usern...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>000103f0d9cfb60f</td>\n",
              "      <td>D'aww! He matches this background colour I'm s...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>000113f07ec002fd</td>\n",
              "      <td>Hey man, I'm really not trying to edit war. It...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>0001b41b1c6bb37e</td>\n",
              "      <td>\"\\nMore\\nI can't make any real suggestions on ...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>0001d958c54c6e35</td>\n",
              "      <td>You, sir, are my hero. Any chance you remember...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-3f2b0022-91d6-465b-b0c9-128092b45884')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-3f2b0022-91d6-465b-b0c9-128092b45884 button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-3f2b0022-91d6-465b-b0c9-128092b45884');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ],
            "text/plain": [
              "                 id  ... identity_hate\n",
              "0  0000997932d777bf  ...             0\n",
              "1  000103f0d9cfb60f  ...             0\n",
              "2  000113f07ec002fd  ...             0\n",
              "3  0001b41b1c6bb37e  ...             0\n",
              "4  0001d958c54c6e35  ...             0\n",
              "\n",
              "[5 rows x 8 columns]"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bDVf_wTaGgVP"
      },
      "source": [
        "train_df, val_df = train_test_split(df, test_size=0.15)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3ro0J6WPGv69",
        "outputId": "d500572e-b917-410e-b791-08302552aa10"
      },
      "source": [
        "# subsample clean comments\n",
        "LABEL_COLUMNS = df.columns.tolist()[2:]\n",
        "\n",
        "train_toxic = train_df[train_df[LABEL_COLUMNS].sum(axis=1) > 0]\n",
        "train_clean = train_df[train_df[LABEL_COLUMNS].sum(axis=1) == 0]\n",
        "\n",
        "train_df = pd.concat([\n",
        "  train_toxic,\n",
        "  train_clean.sample(15_000)\n",
        "])\n",
        "\n",
        "train_df.shape, val_df.shape"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((28754, 8), (23936, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EF3tWEZ74Vww",
        "outputId": "8a0349ae-2010-44b7-a594-bf97d6e95ebd"
      },
      "source": [
        "# take only a subsample of each train_df and val_df for faster iterations\n",
        "train_df = train_df.sample(1000)\n",
        "val_df = val_df.sample(1000)\n",
        "\n",
        "train_df.shape, val_df.shape"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((1000, 8), (1000, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dn-Q0BjUoB9Y"
      },
      "source": [
        "## Creating Dataset and Lightning Data Module"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QV-deDgoo1Rx"
      },
      "source": [
        "# set batch size and max seq_len\n",
        "bs = 12\n",
        "seq_len = 256"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RL6vWIO4HoO2"
      },
      "source": [
        "class ToxicCommentsDataset(Dataset):\n",
        "\n",
        "  def __init__(\n",
        "    self,\n",
        "    data: pd.DataFrame,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    max_token_len: int = 128\n",
        "\n",
        "  ):\n",
        "    self.tokenizer = tokenizer\n",
        "    self.data = data\n",
        "    self.max_token_len = max_token_len\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.data)\n",
        "\n",
        "  def __getitem__(self, index: int):\n",
        "    data_row = self.data.iloc[index]\n",
        "    comment_text = data_row.comment_text\n",
        "    labels = data_row[LABEL_COLUMNS]\n",
        "\n",
        "    encoding = self.tokenizer.encode_plus(\n",
        "      comment_text,\n",
        "      add_special_tokens=True,\n",
        "      max_length=self.max_token_len,\n",
        "      return_token_type_ids=False,\n",
        "      padding=\"max_length\",\n",
        "      truncation=True,\n",
        "      return_attention_mask=True,\n",
        "      return_tensors='pt',\n",
        "\n",
        "    )\n",
        "\n",
        "    return dict(\n",
        "      # comment_text=comment_text, # don't put text here\n",
        "      input_ids=encoding[\"input_ids\"].flatten(),\n",
        "      attention_mask=encoding[\"attention_mask\"].flatten(),\n",
        "      labels = torch.IntTensor(labels)\n",
        "    #   labels=torch.FloatTensor(labels)\n",
        "    )"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mMPUJ8inH2Jg",
        "outputId": "c2ccd0a2-47f0-4908-fab8-9f3d0cdc1087"
      },
      "source": [
        "# test\n",
        "train_dataset = ToxicCommentsDataset(\n",
        "  train_df,\n",
        "  tokenizer,\n",
        "  max_token_len=seq_len\n",
        ")\n",
        "\n",
        "sample_item = train_dataset[0]\n",
        "sample_item.keys()"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['input_ids', 'attention_mask', 'labels'])"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vvh5OltLILMa",
        "outputId": "19bc6636-bde5-4a9d-bab9-201b2feffb2e"
      },
      "source": [
        "print(sample_item[\"input_ids\"], sample_item[\"labels\"])\n",
        "print(sample_item[\"input_ids\"].shape)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([    0,  9178,   171,   481,  3156,    33,    47, 32840,   116,    99,\n",
            "           16,   110,    92,  3018,   766,    35,   741, 14141, 14141,  1988,\n",
            "         1120,   116,     2,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1]) tensor([0, 0, 0, 0, 1, 0], dtype=torch.int32)\n",
            "torch.Size([256])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k0gIaph4IkSu"
      },
      "source": [
        "class ToxicCommentsDataModule(pl.LightningDataModule):\n",
        "  def __init__(self, train_df, test_df, tokenizer, batch_size=8, max_token_len=512):\n",
        "    super().__init__()\n",
        "    self.train_df, self.test_df = train_df, test_df\n",
        "    self.tokenizer = tokenizer\n",
        "    self.batch_size = batch_size\n",
        "    self.max_token_len = max_token_len\n",
        "\n",
        "  def setup(self, stage=None):\n",
        "    self.train_dataset = ToxicCommentsDataset(\n",
        "        self.train_df,\n",
        "        self.tokenizer,\n",
        "        self.max_token_len\n",
        "    )\n",
        "\n",
        "    self.test_dataset = ToxicCommentsDataset(\n",
        "        self.test_df,\n",
        "        self.tokenizer,\n",
        "        self.max_token_len\n",
        "    )\n",
        "\n",
        "  def train_dataloader(self):\n",
        "    return DataLoader(\n",
        "        self.train_dataset,\n",
        "        batch_size = self.batch_size,\n",
        "        shuffle = True,\n",
        "        num_workers=1\n",
        "    )\n",
        "\n",
        "  def val_dataloader(self):\n",
        "    return DataLoader(\n",
        "        self.test_dataset,\n",
        "        batch_size = self.batch_size,\n",
        "        shuffle = False,\n",
        "        num_workers=1\n",
        "    )\n",
        "\n",
        "  def test_dataloader(self):\n",
        "    return DataLoader(\n",
        "        self.test_dataset,\n",
        "        batch_size = self.batch_size,\n",
        "        shuffle = False,\n",
        "        num_workers=1\n",
        "    )"
      ],
      "execution_count": 86,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fwxqBy-FIjf5"
      },
      "source": [
        "data_module = ToxicCommentsDataModule(\n",
        "    train_df,\n",
        "    val_df,\n",
        "    tokenizer,\n",
        "    batch_size = bs,\n",
        "    max_token_len = seq_len\n",
        ")\n"
      ],
      "execution_count": 87,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data_module.setup() # call this before getting len of dataloader\n",
        "len(data_module.train_dataloader())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZtuF0ESKrBv-",
        "outputId": "63cdc40d-6101-4196-d177-98fa759c8289"
      },
      "execution_count": 89,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "84"
            ]
          },
          "metadata": {},
          "execution_count": 89
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mwG-to8c4pUV"
      },
      "source": [
        "# SAM"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6pVZgHoG4qVE"
      },
      "source": [
        "class SAM(torch.optim.Optimizer):\n",
        "    def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):\n",
        "        assert rho >= 0.0, f\"Invalid rho, should be non-negative: {rho}\"\n",
        "\n",
        "        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)\n",
        "        super(SAM, self).__init__(params, defaults)\n",
        "\n",
        "        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)\n",
        "        self.param_groups = self.base_optimizer.param_groups\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def first_step(self, zero_grad=False):\n",
        "        grad_norm = self._grad_norm()\n",
        "        for group in self.param_groups:\n",
        "            scale = group[\"rho\"] / (grad_norm + 1e-12)\n",
        "\n",
        "            for p in group[\"params\"]:\n",
        "                if p.grad is None: continue\n",
        "                self.state[p][\"old_p\"] = p.data.clone()\n",
        "                e_w = (torch.pow(p, 2) if group[\"adaptive\"] else 1.0) * p.grad * scale.to(p)\n",
        "                p.add_(e_w)  # climb to the local maximum \"w + e(w)\"\n",
        "\n",
        "        if zero_grad: self.zero_grad()\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def second_step(self, zero_grad=False):\n",
        "        for group in self.param_groups:\n",
        "            for p in group[\"params\"]:\n",
        "                if p.grad is None: continue\n",
        "                p.data = self.state[p][\"old_p\"]  # get back to \"w\" from \"w + e(w)\"\n",
        "\n",
        "        self.base_optimizer.step()  # do the actual \"sharpness-aware\" update\n",
        "\n",
        "        if zero_grad: self.zero_grad()\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def step(self, closure=None):\n",
        "        assert closure is not None, \"Sharpness Aware Minimization requires closure, but it was not provided\"\n",
        "        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass\n",
        "\n",
        "        self.first_step(zero_grad=True)\n",
        "        closure()\n",
        "        self.second_step()\n",
        "\n",
        "    def _grad_norm(self):\n",
        "        shared_device = self.param_groups[0][\"params\"][0].device  # put everything on the same device, in case of model parallelism\n",
        "        norm = torch.norm(\n",
        "                    torch.stack([\n",
        "                        ((torch.abs(p) if group[\"adaptive\"] else 1.0) * p.grad).norm(p=2).to(shared_device)\n",
        "                        for group in self.param_groups for p in group[\"params\"]\n",
        "                        if p.grad is not None\n",
        "                    ]),\n",
        "                    p=2\n",
        "               )\n",
        "        return norm\n",
        "\n",
        "    def load_state_dict(self, state_dict):\n",
        "        super().load_state_dict(state_dict)\n",
        "        self.base_optimizer.param_groups = self.param_groups"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D1zQOd2AJl7l"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a99EheTsnJyK",
        "outputId": "11f68d95-c98c-4848-f91d-23bcb8efb2fe"
      },
      "source": [
        "sample_batch = next(iter(DataLoader(train_dataset, batch_size=bs, num_workers=1)))\n",
        "sample_batch[\"input_ids\"].shape, sample_batch[\"attention_mask\"].shape"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([12, 256]), torch.Size([12, 256]))"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u8LOwh55Okdz"
      },
      "source": [
        "# test lr schedule with https://www.kaggle.com/isbhargav/guide-to-pytorch-learning-rate-scheduling\n",
        "# test wtih flat cos lr"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ad0Y-3mUXmRM"
      },
      "source": [
        "class BertModel(nn.Module):\n",
        "  def __init__(self, n_classes: int):\n",
        "    super().__init__()\n",
        "    self.model = model\n",
        "    self.classifier = nn.Linear(self.model.config.hidden_size, n_classes)\n",
        "\n",
        "  def forward(self, input_ids, attention_mask): \n",
        "    out = self.model(input_ids, attention_mask=attention_mask)\n",
        "    out = self.classifier(out.pooler_output)\n",
        "    return out"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Eb5FtVEYWQQ"
      },
      "source": [
        "bert_model = BertModel(len(LABEL_COLUMNS)).to(\"cuda\")\n",
        "\n"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uc6_SAgKYmmC"
      },
      "source": [
        "# bert_model(sample_batch[\"input_ids\"], sample_batch[\"attention_mask\"]).shape  # should be bs x 6"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# model saver to save only pytorch model"
      ],
      "metadata": {
        "id": "ea-DIDALncgk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ModelSaver:\n",
        "  def __init__(self, save_path: str, mode: str = \"max\"):\n",
        "    \"\"\"\n",
        "    class used for saving models during training\n",
        "    \"\"\"\n",
        "    self.save_path = save_path\n",
        "    self.mode = mode\n",
        "    assert self.mode in [\"min\", \"max\"], f\"mode {mode} not found\"\n",
        "    # self.best_value = torch.tensor(float(\"inf\")) if mode == \"min\" else torch.tensor(float(\"-inf\"))\n",
        "    self.best_value = float(\"inf\") if mode == \"min\" else float(\"-inf\")\n",
        "\n",
        "  def save_model(self, epoch: int, model: nn.Module, current_value: float):\n",
        "    \"\"\"\n",
        "    compares current_value with self.best_value\n",
        "    if current_value is better then\n",
        "    1. save model weights to self.save_path\n",
        "    2. update self.best_value with current_value\n",
        "    \"\"\"\n",
        "    if (\n",
        "      (self.mode == \"min\" and current_value <= self.best_value)\n",
        "      or (self.mode == \"max\" and current_value >= self.best_value)\n",
        "    ):\n",
        "      torch.save(model.state_dict(), self.save_path)\n",
        "      print(f\"better model found at epoch {epoch} with value {current_value}\")\n",
        "      print(f\"model weights saved to '{self.save_path}', to load model weights, create new model and use new_model.load_state_dict(torch.load('{self.save_path}'))\")\n",
        "      self.best_value = current_value"
      ],
      "metadata": {
        "id": "AejL2yObnfF3"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# PL Module"
      ],
      "metadata": {
        "id": "nQuycz1_nbQR"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XgaReoIsJh0M"
      },
      "source": [
        "class ToxicCommentClassifier(pl.LightningModule):\n",
        "  def __init__(\n",
        "      self, \n",
        "      pytorch_model: nn.Module, \n",
        "      total_steps: int, \n",
        "      lr: float = 2e-5, \n",
        "      rho: float = 0.05,\n",
        "      asam: bool = False,\n",
        "      fit_func: str = \"one_cycle\", \n",
        "      is_ddp: bool = False,\n",
        "      use_ema: bool = False,\n",
        "      accumulate_grad_batches: int = 1,\n",
        "      model_saver: Optional[ModelSaver] = None\n",
        "               \n",
        "  ):\n",
        "    super().__init__()\n",
        "    self.lr = lr\n",
        "    self.rho = rho  # neighborhood size for SAM\n",
        "    self.asam = asam  # whether to use adaptive sam\n",
        "    self.total_steps = total_steps # for lr schedule\n",
        "    self.fit_func = fit_func\n",
        "    self.is_ddp = is_ddp \n",
        "    self.use_ema = use_ema\n",
        "    self.ema = None\n",
        "    self.grad_acc_batches = accumulate_grad_batches\n",
        "    \n",
        "    self.pytorch_model = pytorch_model\n",
        "    self.criterion = nn.BCEWithLogitsLoss() # bce for multi-label\n",
        "\n",
        "    if self.use_ema:\n",
        "      # self.pytorch_model should be moved to device\n",
        "      self.ema = ExponentialMovingAverage(self.pytorch_model.parameters(), decay=0.995)\n",
        "\n",
        "    # metrics \n",
        "    self.f1 = torchmetrics.F1()\n",
        "\n",
        "    # manually define opt step\n",
        "    self.automatic_optimization = False\n",
        "\n",
        "    # track lr schedule\n",
        "    self.lr_schedule = []\n",
        "\n",
        "    # keep track of inputs for sam 2nd update with gradient accumulation\n",
        "    self.input_list, self.attn_mask_list, self.labels_list = [], [], []\n",
        "\n",
        "    # model saving logic\n",
        "    self.model_saver = model_saver\n",
        "\n",
        "  def on_fit_start(self):\n",
        "    \"\"\"\n",
        "    at the beginning of training\n",
        "    - set up scaler for mixed precision (used for SAM manual updates)\n",
        "    \"\"\"\n",
        "    # track scaler for manual optimization\n",
        "    self.scaler = None\n",
        "    if hasattr(self.trainer.precision_plugin, \"scaler\"):\n",
        "      self.scaler = self.trainer.precision_plugin.scaler\n",
        "    print(self.scaler)\n",
        "\n",
        "    print(self.trainer.model)\n",
        "\n",
        "  def forward(self, input_ids, attention_mask):\n",
        "    \"\"\"\n",
        "    forward step for lightning module\n",
        "    \"\"\"\n",
        "    out = self.pytorch_model(input_ids, attention_mask=attention_mask)\n",
        "    return out\n",
        "\n",
        "  def training_step(self, batch, batch_idx):\n",
        "    \"\"\"\n",
        "    training step\n",
        "    since using manual optimization: only precision and accelerator logic is handled by Lightning\n",
        "\n",
        "    Must manually specify everything else, such as:\n",
        "    - manual optimization logic ie \n",
        "      - opt = self.optimizers()\n",
        "      - opt.zero_grad()\n",
        "      - self.manual_backward(loss)\n",
        "      - and opt.step()\n",
        "    - learning rate scheduler step\n",
        "    - gradient accumulation logic\n",
        "    - gradient clipping\n",
        "    - etc\n",
        "    \"\"\"\n",
        "    # batch comes from dataset\n",
        "    input_ids = batch[\"input_ids\"]\n",
        "    attention_mask = batch[\"attention_mask\"]\n",
        "    labels = batch[\"labels\"]\n",
        "\n",
        "    # save input and output for 2nd step\n",
        "    self.input_list.append(input_ids)\n",
        "    self.attn_mask_list.append(attention_mask)\n",
        "    self.labels_list.append(labels)\n",
        "\n",
        "\n",
        "    out = self(input_ids, attention_mask)\n",
        "    loss = self.criterion(out, labels.to(dtype=torch.float32)) # cast to float because labels is in Int\n",
        "    # scale the loss by gradient accumulation batches\n",
        "    loss = loss / self.grad_acc_batches\n",
        "\n",
        "    # track f1, use sigmoid first\n",
        "    f1 = self.f1(torch.sigmoid(out), labels)\n",
        "\n",
        "    \"\"\"\n",
        "    IMPORTANT NOTE:\n",
        "    need to manually handle scaling with optimizer\n",
        "    need to call opt.optimizer.first_step_step(zero_grad=True)\n",
        "    and manaually scale/unscale gradients and call scaler.update()\n",
        "    \"\"\"\n",
        "\n",
        "    # loss backward\n",
        "    if self.is_ddp:\n",
        "      # maybe should be `with self.pytorch_model.no_sync()`?\n",
        "      # (test with both on ddp)\n",
        "      with self.trainer.model.no_sync():\n",
        "        self.manual_backward(loss)\n",
        "    else:\n",
        "      self.manual_backward(loss)\n",
        "\n",
        "    # optimizer step every grad_acc_batches steps\n",
        "    if (batch_idx + 1) % self.grad_acc_batches == 0:\n",
        "      \"\"\"\n",
        "      must use manual optimization with scaling logic eg\n",
        "      https://github.com/PyTorchLightning/pytorch-lightning/discussions/11290\n",
        "\n",
        "      ```\n",
        "      optimizer = self.optimizers().optimizer # access pytorch optimizer from lightningoptimizer\n",
        "      scaler = self.trainer.precision_plugin.scalar\n",
        "      if scaler:\n",
        "        scaler.unscale_(optimizer)\n",
        "        optimizer.first_step(zero_grad=True)\n",
        "        scaler.update()\n",
        "\n",
        "      else:\n",
        "        optimizer.first_step(zero_grad=True)\n",
        "\n",
        "      # 2nd pass\n",
        "      ...\n",
        "\n",
        "      ```\n",
        "      \"\"\"\n",
        "      # acecss the optimizer\n",
        "      # since using sam, want to use the base pytorch optimizer\n",
        "      # get the pytorch optimizer from lightning optimizer\n",
        "      optimizer = self.optimizers().optimizer \n",
        "      self.lr_schedule.append(optimizer.param_groups[0][\"lr\"]) # check that lrs are the same (is true)\n",
        "\n",
        "      if self.scaler:\n",
        "        self.scaler.unscale_(optimizer)\n",
        "        optimizer.first_step(zero_grad=True)\n",
        "        self.scaler.update()\n",
        "      else:\n",
        "        optimizer.first_step(zero_grad=True)\n",
        "\n",
        "      # 2nd forward pass with saved input_list and labels_list\n",
        "      # to get the accumulated gradients again\n",
        "      for (input_ids, attention_mask, labels) in list(zip(self.input_list, self.attn_mask_list, self.labels_list)):\n",
        "        out_2 = self(input_ids, attention_mask)\n",
        "        loss_2 = self.criterion(out_2, labels.to(dtype=torch.float32))\n",
        "        loss_2 = loss_2 / self.grad_acc_batches\n",
        "        self.manual_backward(loss_2)\n",
        "        # 2nd optimizer step\n",
        "        if self.scaler:\n",
        "          self.scaler.unscale_(optimizer)\n",
        "          optimizer.second_step(zero_grad=True)\n",
        "          self.scaler.update()\n",
        "        else:\n",
        "          optimizer.second_step(zero_grad=True)\n",
        "\n",
        "      # clear saved lists\n",
        "      self.input_list, self.attn_mask_list, self.labels_list = [], [], []\n",
        "\n",
        "      # update ema (after every optimizer step)\n",
        "      if self.use_ema:\n",
        "        self.ema.update() # note: this might be self.ema.update(self.trainer.model.parameters()) or maybe self.parameters()?\n",
        "        # self.ema.update(self.trainer.model.parameters())\n",
        "\n",
        "      # lr schedule (every optimizer step)\n",
        "      lr_sch = self.lr_schedulers()\n",
        "      lr_sch.step()\n",
        "\n",
        "    # log loss and metrics each k steps and each epoch\n",
        "    loss = loss * self.grad_acc_batches  # scale loss back wrt gradient accumulation\n",
        "    log_values = {\"loss\": loss, \"train_f1\": f1}\n",
        "    self.log_dict(log_values, sync_dist=True, prog_bar=True, on_step=True, on_epoch=False)\n",
        "    return {\"loss\": loss, \"train_f1\": f1}\n",
        "\n",
        "  def validation_step(self, batch, batch_idx):\n",
        "    \"\"\"\n",
        "    validation step\n",
        "    get validation loss and accuracy metrics\n",
        "    \"\"\"\n",
        "    # batch comes from dataset\n",
        "    input_ids = batch[\"input_ids\"]\n",
        "    attention_mask = batch[\"attention_mask\"]\n",
        "    labels = batch[\"labels\"]\n",
        "\n",
        "    if self.use_ema:\n",
        "      with self.ema.average_parameters():\n",
        "        out = self(input_ids, attention_mask)\n",
        "    else:\n",
        "      out = self(input_ids, attention_mask)\n",
        "\n",
        "    # compute loss + metrics\n",
        "    loss = self.criterion(out, labels.to(dtype=torch.float32)) # cast to float because labels is in Int\n",
        "    f1 = self.f1(out, labels)\n",
        "    # log loss and metrics each k steps and each epoch\n",
        "    log_values = {\"val_loss\": loss, \"val_f1\": f1}\n",
        "    self.log_dict(log_values, sync_dist=True, prog_bar=True, on_step=True, on_epoch=False)\n",
        "    return {\"val_loss\": loss, \"val_f1\": f1}\n",
        "\n",
        "  def training_epoch_end(self, outputs):\n",
        "    \"\"\"\n",
        "    after every training epoch, log train metrics\n",
        "\n",
        "    note: this might not be needed if logging is done for each step and epoch\n",
        "    \"\"\"\n",
        "    # log validation loss and metrics\n",
        "    avg_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
        "    avg_f1 = torch.stack([x[\"train_f1\"] for x in outputs]).mean()\n",
        "\n",
        "    self.log(\"avg_train_loss\", avg_loss, sync_dist=True, prog_bar=True)\n",
        "    self.log(\"avg_train_f1\", avg_f1, sync_dist=True, prog_bar=True)\n",
        "\n",
        "  def validation_epoch_end(self, outputs):\n",
        "    \"\"\"\n",
        "    after every training epoch, log val metrics\n",
        "    also save pytorch model weights to be loaded and decoupled\n",
        "    \"\"\"\n",
        "    # log validation loss and metrics\n",
        "    avg_loss = torch.stack([x[\"val_loss\"] for x in outputs]).mean()\n",
        "    avg_f1 = torch.stack([x[\"val_f1\"] for x in outputs]).mean()\n",
        "    self.log(\"avg_val_loss\", avg_loss, sync_dist=True, prog_bar=True)\n",
        "    self.log(\"avg_val_f1\", avg_f1, sync_dist=True, prog_bar=True)\n",
        "\n",
        "    if self.model_saver:\n",
        "      if self.use_ema:\n",
        "        with self.ema.average_parameters():\n",
        "          self.model_saver.save_model(self.current_epoch, self.pytorch_model, avg_f1)\n",
        "      else:\n",
        "        self.model_saver.save_model(self.current_epoch, self.pytorch_model, avg_f1)\n",
        "\n",
        "  def configure_optimizers(self):\n",
        "    # TODO: \n",
        "    # 1. automatically find steps for onecycleLR\n",
        "    # 2. add flat cos\n",
        "    # 3. make lr scheduler configurable(one_cycle vs flat_cos)\n",
        "\n",
        "    # trying SAM\n",
        "    base_optimizer = AdamW\n",
        "    optimizer = SAM(\n",
        "        self.pytorch_model.parameters(), \n",
        "        base_optimizer=base_optimizer, \n",
        "        lr=self.lr, \n",
        "        betas=(0.9, 0.99),\n",
        "        rho=self.rho,\n",
        "        adaptive=self.asam\n",
        "    )\n",
        "    # lr schedule with sam (seems to work fine with 2 step?)\n",
        "    scheduler = OneCycleLR(\n",
        "      optimizer=optimizer,\n",
        "      max_lr=self.lr,\n",
        "      pct_start=0.3,\n",
        "      total_steps=self.total_steps\n",
        "    )\n",
        "    \n",
        "    return [optimizer], [scheduler]\n",
        "\n",
        "  def on_save_checkpoint(self, checkpoint):\n",
        "    \"\"\"\n",
        "    saving model weights, configurable with EMA\n",
        "    \"\"\"\n",
        "    if self.ema:\n",
        "      with self.ema.average_parameters():\n",
        "        checkpoint['state_dict'] = self.state_dict()\n",
        "    else:\n",
        "      checkpoint['state_dict'] = self.state_dict()"
      ],
      "execution_count": 106,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHhFIJkGUTU4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a4cbe8cc-636e-4c5e-f0bd-05830c99bf64"
      },
      "source": [
        "# training epoch related hyperparams\n",
        "steps_per_epoch = len(data_module.train_dataloader())\n",
        "n_epochs = 1\n",
        "accumulate_grad_batches = 4\n",
        "n_gpus = 1\n",
        "\n",
        "\n",
        "total_steps = int(steps_per_epoch * n_epochs / n_gpus / accumulate_grad_batches)\n",
        "print(f\"total_steps: {total_steps}\")\n",
        "\n",
        "\n",
        "# creating model saver\n",
        "model_saver = ModelSaver(save_path=\"test.pth\", mode=\"max\")"
      ],
      "execution_count": 107,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total_steps: 21\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPhpGLWyV0eL"
      },
      "source": [
        "Note: if using multiple gpus, total steps calculation is something like\n",
        "```\n",
        "steps_per_epoch = len(data_module.train_dataloader())\n",
        "\n",
        "# for gpu or tpu cores\n",
        "num_devices = 4\n",
        "if tpu_cores:\n",
        "    num_devices = max(num_devices, tpu_cores)\n",
        "\n",
        "accumulate_grad_batches = 1 # for no gradient accumulation\n",
        "\n",
        "effective_accum = accumulate_grad_batches * num_devices\n",
        "total_steps (steps_per_epoch // effective_accum) * n_epochs\n",
        "\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "del bert_model\n",
        "bert_model = BertModel(len(LABEL_COLUMNS)).to(\"cuda\")\n"
      ],
      "metadata": {
        "id": "JCbxCF88NrTO"
      },
      "execution_count": 108,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O5o32pPDKTTi"
      },
      "source": [
        "toxic_comment_model = ToxicCommentClassifier(\n",
        "  bert_model, \n",
        "  total_steps = total_steps, \n",
        "  use_ema=True, \n",
        "  lr=0.1, \n",
        "  accumulate_grad_batches=accumulate_grad_batches,\n",
        "  model_saver=model_saver\n",
        "\n",
        ")"
      ],
      "execution_count": 109,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fWWO7UHn_Plj",
        "outputId": "a7823802-d38f-407b-fd6d-ea8acc7f16c7"
      },
      "source": [
        "sam = toxic_comment_model.configure_optimizers()[0][0]\n",
        "sam.param_groups[0][\"lr\"]"
      ],
      "execution_count": 110,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0040000000000000036"
            ]
          },
          "metadata": {},
          "execution_count": 110
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "li0_cLa82LXx"
      },
      "source": [
        "# Trainer\n",
        "\n",
        "defining the trainer and running `trainer.fit` "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XrBKqeKnKWy-"
      },
      "source": [
        "# callbacks\n",
        "lr_monitor_cb = LearningRateMonitor(logging_interval='step')\n",
        "\n",
        "checkpoint_callback = ModelCheckpoint(\n",
        "  dirpath=\"checkpoints\",\n",
        "  filename=\"best-checkpoint\",\n",
        "  save_top_k=1,\n",
        "  verbose=True,\n",
        "  monitor=\"avg_val_f1\",\n",
        "  mode=\"max\"\n",
        ")\n",
        "\n",
        "early_stopping_callback = EarlyStopping(monitor='avg_val_loss', patience=6)\n",
        "\n",
        "logger = TensorBoardLogger(\"lightning_logs\", name=\"toxic-comments\")\n",
        "\n",
        "# standard callback, checkpoint and early stopping\n",
        "# callbacks = [lr_monitor_cb, checkpoint_callback, early_stopping_callback]\n",
        "\n",
        "callbacks = [early_stopping_callback]"
      ],
      "execution_count": 111,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cZdITHvZhBSF",
        "outputId": "e8ffc8a6-d156-443b-e52e-7a55643f78f6"
      },
      "source": [
        "trainer = pl.Trainer(\n",
        "    logger=logger,\n",
        "    callbacks = callbacks,\n",
        "    precision=32,\n",
        "    max_epochs = n_epochs,\n",
        "    accelerator=\"gpu\",\n",
        "    devices=n_gpus,\n",
        "    num_sanity_val_steps=0\n",
        ")"
      ],
      "execution_count": 112,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "GPU available: True, used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FSIidqFNT5n5"
      },
      "source": [
        "# optional: currently not working?\n",
        "# lr_finder = trainer.tuner.lr_find(model=toxic_comment_model, datamodule=data_module)"
      ],
      "execution_count": 113,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jfl3uth_hSRg",
        "outputId": "c6c9e024-dc70-4dc6-daa9-d27fe72dcc79"
      },
      "source": [
        "trainer.fit(toxic_comment_model, data_module)"
      ],
      "execution_count": 114,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py:686: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...\n",
            "  rank_zero_warn(\"Detected KeyboardInterrupt, attempting graceful shutdown...\")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AtRPI8ocIRai",
        "outputId": "620da822-6976-4585-8b43-596bf4ea81e1"
      },
      "source": [
        "trainer.logged_metrics"
      ],
      "execution_count": 100,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'avg_train_f1': 0.16689248383045197,\n",
              " 'avg_train_loss': 1.9315109252929688,\n",
              " 'avg_val_f1': 0.0,\n",
              " 'avg_val_loss': 0.3577773869037628,\n",
              " 'loss': tensor(1.7774),\n",
              " 'train_f1': tensor(0.),\n",
              " 'val_f1': 0.0,\n",
              " 'val_loss': 0.07361804693937302}"
            ]
          },
          "metadata": {},
          "execution_count": 100
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# validate lr scheduler is working correctly\n",
        "print(toxic_comment_model.lr_schedule)\n",
        "print(len(toxic_comment_model.lr_schedule), total_steps)\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "plt.plot(toxic_comment_model.lr_schedule)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 340
        },
        "id": "4UPsAtlFgIG5",
        "outputId": "c7fbc6af-caf4-4ba8-e3b3-44575008eb61"
      },
      "execution_count": 101,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.0040000000000000036, 0.012188527098390314, 0.03396027605851658, 0.06188697369957272, 0.08644034833812791, 0.09924306774148466, 0.09944154354509119, 0.09673623556045154, 0.09190443762697105, 0.08516599658482699, 0.07682751154913771, 0.0672683836480545, 0.056923553268744854, 0.046263710266697504, 0.035773877573965704, 0.02593134265232537, 0.01718394091038552, 0.009929679189326548, 0.004498626448808909, 0.0011378956269720922, 4e-07]\n",
            "21 21\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[<matplotlib.lines.Line2D at 0x7f4b595c03d0>]"
            ]
          },
          "metadata": {},
          "execution_count": 101
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXoAAAD4CAYAAADiry33AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO3de1yUZf7/8ddnOIsICAoKIh5T8AweO26lWWtiZy0322rtsLW11a9v7X632trarbat7WRZdrIts6NWlnbG0kwUPIAHEI8oCKgoIufr98eMfVlCHWRm7pnh83w8fDAz9z3O25vhzXjd933dYoxBKaWU/7JZHUAppZR7adErpZSf06JXSik/p0WvlFJ+ToteKaX8XKDVAZqLjY01ycnJVsdQSimfsmrVqjJjTJeWlnld0ScnJ5OVlWV1DKWU8ikisv1Yy3ToRiml/JwWvVJK+TkteqWU8nNa9Eop5eecKnoRmSgim0SkQETuaWH5GSKyWkTqReTSZstmiEi+488MVwVXSinlnBMWvYgEAM8B5wMpwDQRSWm22g7gGuCtZs/tDNwPjAZGAfeLSHTbYyullHKWM5/oRwEFxphCY0wtMA/IaLqCMWabMWYt0NjsuecBXxhj9hlj9gNfABNdkFsppZSTnDmOPgHY2eT+Luyf0J3R0nMTmq8kIjOBmQBJSUlO/tWqrQ5W17Gvspa6hkZq6hupbWikzvG1tr7x58frGgy19Y3U1jfYbzuWD4iP4FcDuhIaFGD1P0UpdRxeccKUMWY2MBsgPT1dJ8j3gI3FB7nk+WUcrm1o098TERLIeYPimTIsgbF9YgiwiYsSKqVcxZmiLwJ6NLmf6HjMGUXAWc2e+62Tz1VuYozhvo9yCQ608deMQYQE2ggKsBESaCPYcTs40EZw86+BNoIChOBAGwEi/Fi4j49yivh8fTHvrdpFl4gQLhzSnYxh3RmSGImIlr5S3kBOdIUpEQkENgPnYC/ulcCVxpjcFtZ9DfjEGPOe435nYBUwwrHKaiDNGLPvWK+Xnp5udAoE9/oou4jb38nh7xcPZtqotg+VVdc18PXGvSzIKeKbjaXUNjTSKzacyUO7M2V4Ar1iw12QWil1PCKyyhiT3uIyZy4lKCIXAE8BAcArxpiHReRBIMsYs1BERgIfAtFANVBsjEl1PPda4E+Ov+phY8yrx3stLXr3OlRdx9lPfEf3yFA+vPlUbC4eaqmoquPz3D18lL2bH7eWYwwMSYwkY1gCFw7pRtdOoS59PaWUXZuL3pO06N3rb5/kMeeHrXx086kM7RHl1tcqrqjm4zW7WbCmiPVFB7EJjOsTy+Rh3Tl/UDwRoUFufX2l2hMtegXA5pJDnP/vpVyensjfLx7i0dcu2HuIhTm7WbBmN9vLq4juEMRfJqVw0fAEHctXygW06BXGGKa99CMbiw/x9Z1n0Tk82LIcq3fs5+FPN7B6xwFO7RvDw1MGk6zj+Eq1yfGKXue6aSc+XruHHwv3cdeEUywreQARIa1nZ967cRwPTRnE2p0VnPdUJs99U0BdQ/Pz7ZRSrqBF3w5U1tTz8Kd5DEro5JKjbFzBZhN+M6YnX955JmcP6Mrjizdx4TPfs3rHfqujKeV3tOjbgWe+yqfkYA0PZgzyuhOa4jqFMmt6Gi9dnU7FkToumbWM+xas51B1ndXRlPIbWvR+Lr/kEHO+38oV6T0YkeS988mNT4njizvOZMbYZOb+uJ3x/8pkcW6x1bGU8gta9H7MGMP9C3PpEBzA3RNPsTrOCXUMCeSByal8ePOpRIcHc8PcVcx8I4s9FUesjqaUT9Oi92OfrtvDsi3l/L/zTiGmY4jVcZw2rEcUC285lXvOH0Bmfinj/5XJG8u30dDoXUeIKeUrtOj91OGaev72yQZSu3fiytE9rY7TakEBNm48sw9Lbj+T4UlR3Lcgl0tmLWPDnoNWR1PK52jR+6lnvi6g+GC1V+6AbY2kmA68ce0onrpiGDv3VXHhM9/z+OKN1OuhmEo5TYveD20prWTO94VcmpZIWk/v3QHrLBFhyvAEvrzjTKYMT+C5b7Zw45urqKqttzqaUj5Bi97PGGN4YGEuoUEB3HP+AKvjuFR0eDD/vGwoD2ak8vXGvUyb/SNllTVWx1LK62nR+5nP1xezNL+MO8f3J9aHdsC2xtVjk3lhehqbSg5x8fPLKCyttDqSUl5Ni96PVNXW89AneQyIj2D6GN/bAdsaE1Ljeft3Y6isqeeSWctYtf2YlzhQqt3Tovcjz31TwO6Kah6aMojAAP//1g5PiuaDm8YRGRbElS+t4PP1eoKVUi3x/zZoJwpLK3kpcysXj0hgZHJnq+N4THJsOO/fNI6U7p246T+reO2HrVZHUsrraNH7AWMMD3ycR0igjXvPH2h1HI+L6RjCW9eP4dyBcTzwcR4Pf5pHo55cpdTPtOj9wOLcEjI3l/LH8f3pEuGfO2BPJCw4gBempzFjbE9eWrqVW+dlU13XYHUspbxCoNUBVNscqW34eQfs1WP9ewfsiQTYhAcmp5IQHcYjizZSerCG2VenEdXBuvn3lfIG+onexz3/bQFFB47w18mp7WIH7ImICDPP6MPT04aTs/MAl8xaxs59VVbHUspS2gw+bEd5FS9+V8iUYd0Z3TvG6jheZfLQ7rxx3ShKD9Vw8axlrC+qsDqSUpbRovdh76/eRX1jI/e0wx2wzhjTO4b3bxpHcICNy19czjeb9lodSSlLaNH7sCV5JaT37Ex8ZKjVUbxWv7gIPrh5HMkx4Vz/ehbvrNxhdSSlPE6L3kft3FfFhj0HGZ8SZ3UUrxfXKZT5N47l1L6x/M/765i7fJvVkZTyKC16H7UkrwRAi95JHUMCmTMjnXMHxnHfwlwWrtltdSSlPEaL3kctyS3mlLgIkmPDrY7iM4ICbDx75XBGJnfmzvk5ZG4utTqSUh6hRe+D9h2uZeW2fUxI1U/zrRUaFMDLM9Lp2zWCG+auInvHfqsjKeV2WvQ+6OuNe2k0MCEl3uooPqlTaBCvXzuSrp1C+O1rK8kvOWR1JKXcSoveBy3JLaZbZCiDEjpZHcVndY0IZe61owkKsHH1Kz9RdOCI1ZGUchsteh9zpLaBzPxSxqfEIeK714L1BkevR1tZU89v5qygXK9WpfyUFr2PWZpfSnVdow7buMjAbp2YM2MkRfuP8NvXVlJZo9ehVf7HqaIXkYkisklECkTknhaWh4jIO47lK0Qk2fF4kIi8LiLrRGSDiNzr2vjtz5K8EiJCAxndu/3MOe9uo3p15vmrRpC7+yA3zM2ipl5nvVT+5YRFLyIBwHPA+UAKME1EUpqtdh2w3xjTF3gSeNTx+GVAiDFmMJAG3HD0l4BqvfqGRr7aUMI5A7oSpBOYudQ5A+N47JIh/FBQzh/fyaFB57NXfsSZthgFFBhjCo0xtcA8IKPZOhnA647b7wHniH0A2QDhIhIIhAG1wEGXJG+HVm3fz/6qOiak6rCNO1ySlsj//nogi9YV85cF6zFGy175B2fmo08Adja5vwsYfax1jDH1IlIBxGAv/QxgD9AB+KMx5hdXcRaRmcBMgKSkpFb+E9qPJXklBAfYOKN/F6uj+K3rT+9N+eFaZn27hZjwYO6ccIrVkZRqM3dfeGQU0AB0B6KBpSLypTGmsOlKxpjZwGyA9PR0/RjVAmMMS/KKObVvDB1D9Hox7nT3eaew/3Atz3xdQOfwYH57ai+rIynVJs4M3RQBPZrcT3Q81uI6jmGaSKAcuBL43BhTZ4zZC/wApLc1dHu0sfgQO/cd0WEbDxARHr5oMBNT4/nrx3l8lN387a6Ub3Gm6FcC/USkl4gEA1OBhc3WWQjMcNy+FPja2Ac4dwBnA4hIODAG2OiK4O3NktwSROCcgV2tjtIuBNiEp6YOY2zvGO56d43OZa982gmL3hhTD9wCLAY2APONMbki8qCITHasNgeIEZEC4A7g6CGYzwEdRSQX+y+MV40xa139j2gPvthQzIikaLpG6NzznhIaFMDsq9MY0C2Cm95cxartv9i9pJRPEG87siA9Pd1kZWVZHcOrFB04wqn/+Jp7zx/ADWf2sTpOu1NWWcPlLyynrLKG924aR/+4CKsjKfULIrLKGNPi0LgejO0DvsgtBnTueavEdgzhjetGERoUwPWvZ3GgqtbqSEq1iha9D1iSV0Lfrh3p3aWj1VHarcToDsyansaeiiPc+na2nlClfIoWvZc7UFXLiq37mKCf5i2X1jOaBzMGsTS/jMcXb7I6jlJO0wOyvdzXG/fS0Gj0sEovMW1UEuuKKnjhuy0MSujEpCHdrY6k1AnpJ3ov90VeCXGdQhiSEGl1FOXwwIWppPWM5v+9u5YNe3RGD+X9tOi9WHVdA99tLuXcgXHYbDr3vLcIDrQx66oRdAoLZOZc3TmrvJ8WvRf7oaCMqtoGHbbxQl07hTJreholFTW6c1Z5PS16L7Ykt4SIkEDG9o6xOopqwYikaB7MSGVpfhmPLdYTvpX30p2xXqqh0fDlhhLOGtCV4ED9feytpjp2zr74XSGDukdy4VDdOau8jzaIl8resZ/yw7V6WKUPuP/CVNJ7RnP3e2vJ2607Z5X30aL3UkvySggKEM46Reee93bBgTaen27fOXvDm1nsP6w7Z5V30aL3QsYYFucWM7ZPLBGhQVbHUU7oGhHKC46ds3+Yl019Q6PVkZT6mRa9F8rfW8n28iodtvExw5OieWhKqp45q7yO7oz1Qkt0EjOfdcVIx87ZzEJSEyKZrDtnlRfQT/ReaEleCcN6RBHXSeee90X3TUplZHI0d7+3RnfOKq+gRe9l9lQcYe2uCiak6qd5XxUcaOO5q0YQFRbMzLm6c1ZZT4vey3yZVwKg4/M+rmtEKLOmj2DvQfuZs7pzVllJi97LLMkroXdsOH107nmfNzwpmr9NGcT3BWU8pjtnlYW06L1IxZE6lm8pZ3xqHCI6iZk/uHxkD34zpiezMwtZkFNkdRzVTmnRe5FvN+2lvtEwIUUnMfMnf5mUwsjkaO79YB2FpZVWx1HtkBa9F1mSW0JsxxCG94iyOopyoeBAG09PG05IoI3fv5VNdV2D1ZFUO6NF7yVq6hv4dtNexqfo3PP+qFtkGE9cPpQNew7yyKINVsdR7YwWvZdYtqWcw7UNerSNHzt7QBzXn9aLN5Zv57N1e6yOo9oRLXovsSS3hPDgAMb20bnn/dndEwcwNDGSu99fy859VVbHUe2EFr0XaGw0fJFXwlmndCU0KMDqOMqNggNtPDNtBBi49e1s6vT4euUBWvReIHvnAcoqa/Rs2HYiKaYD/7hkCDk7D/BPPb5eeYAWvRf4Iq+EQJtw1ildrY6iPOTXQ7px1egkXsws5JtNe62Oo/ycFr0XWJJXzJjeMUSG6dzz7clfJqUwID6CO+evobii2uo4yo9p0VusYG8lhaWHddimHQoNCuDZK0dwpLaB2+Zl09BorI6k/JQWvcWW5Nnnnj93oBZ9e9S3a0cemjKIFVv38fRX+VbHUX7KqaIXkYkisklECkTknhaWh4jIO47lK0QkucmyISKyXERyRWSdiOgk6018t6mUgd060T0qzOooyiKXpiVy8YgEnvk6n+Vbyq2Oo/zQCYteRAKA54DzgRRgmoikNFvtOmC/MaYv8CTwqOO5gcCbwI3GmFTgLKDOZel9XGVNPat37OeM/rFWR1EWeyhjEMmx4dw2L5vyyhqr4yg/48wn+lFAgTGm0BhTC8wDMpqtkwG87rj9HnCO2KdfnACsNcasATDGlBtjdKIPhxWF5dQ1GM7o18XqKMpi4SGBPDttBAeO1HHH/DU06ni9ciFnij4B2Nnk/i7HYy2uY4ypByqAGKA/YERksYisFpG72x7ZfyzNLyM0yEZ6crTVUZQXSOneib9MSuG7zaW8tLTQ6jjKj7h7Z2wgcBpwlePrRSJyTvOVRGSmiGSJSFZpaambI3mPzM2ljOkdQ0igng2r7KaPTuKCwfE8vngTq3fstzqO8hPOFH0R0KPJ/UTHYy2u4xiXjwTKsX/6zzTGlBljqoBFwIjmL2CMmW2MSTfGpHfp0j6GMXbuq6Kw7DCn67CNakJE+PvFQ4iPDOXWt7KpqNJdWqrtnCn6lUA/EeklIsHAVGBhs3UWAjMcty8FvjbGGGAxMFhEOjh+AZwJ5Lkmum/7vqAMgDP66Y5Y9d8iw4J49soRlBys5u7312D/UVLq5J2w6B1j7rdgL+0NwHxjTK6IPCgikx2rzQFiRKQAuAO4x/Hc/cC/sP+yyAFWG2M+df0/w/cszS8lvlMofbvqtWHVLw3rEcX/TBzA4twS5v643eo4yscFOrOSMWYR9mGXpo/d1+R2NXDZMZ77JvZDLJVDQ6Ph+/wyJg6K12vDqmO67rReLC8s52+fbCCtZzSp3SOtjqR8lJ4Za4G1uw5wsLpex+fVcdlswj8vG0rn8GBueSubwzX1VkdSPkqL3gKZm8sQgVP76vi8Or7O4cE8NXUY28oP8+DHuntLnRwtegsszS9lcEIkncODrY6ifMCY3jH8/qy+vJO1k0/X6iUIVetp0XvYweo6snce4HQ92ka1wm3n9mNYjyju/WAtRQeOWB1H+Rgteg9bvqWchkaj4/OqVYICbDw9dTiNBv44L0enNFatokXvYUvzSwkPDmBEkk57oFonKaYDD01J5adt+3j+mwKr4ygfokXvYZmbyxjbJ4bgQN30qvUuGp5IxrDuPPVVPqu26xQJyjnaNh60vfwwO/ZV6bCNapOHpgyiW2Qot83L5mC1TpGgTkyL3oMy8+3THuiOWNUWnUKD+PfU4eypqOa+j9ZbHUf5AC16D1q6uZSEqDB6xYZbHUX5uLSe0dx2Tj8+ytnNh9m7rI6jvJwWvYfUNTSyfEs5Z/TvotMeKJf4/a/6Miq5M3/5KJcd5VVWx1FeTIveQ9bsPMChmnqdrVK5TIBNeHLqMETgD/OyqWtotDqS8lJa9B6SubkUm8C4Plr0ynUSosL4+8WDydl5gKe/yrc6jvJSWvQekplfxtAeUUR2CLI6ivIzk4Z05/L0RJ79poAfC8utjqO8kBa9BxyoqmXtrgN6WKVym/svTCU5Jpw/vpOjV6VSv6BF7wHLtpTTaPRqUsp9wkMC+ffUYZQequHeD9fqVanUf9Gi94Cl+aVEhAQyrEeU1VGUHxuSGMVd553ConXFzM/aaXUc5UW06N3MGEPm5jLG9Y0hMEA3t3Kvmaf3ZlyfGB5YmMeW0kqr4ygvoc3jZoVlhyk6cETH55VH2GzCvy4fRmiQjT+8nU1NfYPVkZQX0KJ3s6WbSwE4Q4teeUh8ZCiPXjKE3N0HeWLJZqvjKC+gRe9mS/PL6BnTgaSYDlZHUe3IhNR4po9JYnZmIUvzS62OoyymRe9GtfWNLC8s10/zyhJ/viCFfl07csf8NZRX1lgdR1lIi96NVu/YT1Vtg85WqSwRFhzA09OGU3GkjrveXUOjXpWq3dKid6Ol+aUE2ISxfWKsjqLaqYHdOvHnCwbyzaZSXl22zeo4yiJa9G6UubmMEUlRRITqtAfKOleP7cn4lDj+8dkG1hdVWB1HWUCL3k3KK2tYv7tCD6tUlhMRHrtkCLEdQ7j17Wwqa+qtjqQ8TIveTX7YUo4xejUp5R2iw4N56ophbC8/zH0L9KpU7Y0WvZss3VxKZFgQQxJ12gPlHUb3juHWs/vxweoivSpVO6NF7wbGGJbml3Fa31gCbHo1KeU9bj3bflWq//1wPdvKDlsdR3mIFr0b5O+tpPhgtQ7bKK8TGGDjqanDCAyw8Yd52dTW61Wp2gMtejfIdEx7cJoWvfJC3aPCeOzSIazdVcHjizdaHUd5gFNFLyITRWSTiBSIyD0tLA8RkXccy1eISHKz5UkiUikid7kmtndbml9G7y7hJEbrtAfKO52XGs9vxvTkpaVb+WbTXqvjKDc7YdGLSADwHHA+kAJME5GUZqtdB+w3xvQFngQebbb8X8BnbY/r/arrGlixVac9UN7vz78eyID4CO6av4a9B6utjqPcyJlP9KOAAmNMoTGmFpgHZDRbJwN43XH7PeAcEREAEZkCbAVyXRPZu63avp/qukbO6K/DNsq7hQYF8My04RyureeO+TpFgj9zpugTgKaXq9nleKzFdYwx9UAFECMiHYH/Af56vBcQkZkikiUiWaWlvj3TXmZ+KUEBwuheOu2B8n794iK4/8JUvi8o48XMQqvjKDdx987YB4AnjTHHvdSNMWa2MSbdGJPepYtvD3lkbi4jrWc04SGBVkdRyilTR/bg14O78cSSTWTv2G91HOUGzhR9EdCjyf1Ex2MtriMigUAkUA6MBh4TkW3A7cCfROSWNmb2WnsPVbNhz0Gd9kD5FBHhkYsHE9cplFvfzuZgdZ3VkZSLOVP0K4F+ItJLRIKBqcDCZussBGY4bl8KfG3sTjfGJBtjkoGngEeMMc+6KLvX+aGgDNCrSSnfExkWxNPThrOnopo/fbAOY3S83p+csOgdY+63AIuBDcB8Y0yuiDwoIpMdq83BPiZfANwB/OIQzPZg6eYyOocHk9q9k9VRlGq1tJ7R3DG+P5+s3cO7WTpFgj9xaiDZGLMIWNTssfua3K4GLjvB3/HASeTzGcYYMh3THth02gPlo248sw8/FJRx/8JcRvSMom/XCKsjKRfQM2NdZGPxIcoqa3TaA+XTAmzCk1cMIyw4gFveyqa6rsHqSMoFtOhd5Oi0B7ojVvm6uE6hPHHZUDYWH+LvizZYHUe5gBa9iyzNL6N/XEfiI0OtjqJUm/1qQFeuO60Xry/fzufri62Oo9pIi94FjtQ28NO2fXq0jfIrd088haGJkdz17hoKS497Kozyclr0LvDTtn3U1jdyen8teuU/QgIDeH56GkEBwk1vrqaqVi9B6Ku06F1g6eZSggNtjErubHUUpVwqISqMp6cNZ/PeQ9yrx9f7LC36NjLG8NXGvYzu1Zmw4ACr4yjlcqf368Kd4/uzIGc3byzfbnUcdRK06NtoXVEFW8sOM2lIN6ujKOU2N5/Vl3MHduVvn+axarvOh+NrtOjbaEHOboIDbExM1aJX/stmE564fBjdIsP4/X9WU1ZZY3Uk1Qpa9G3Q0Gj4eM1uzjqlC5EdgqyOo5RbRYYFMWv6CPZX1XLrW9nUN+j1Zn2FFn0brCgsZ++hGjKGNZ+eXyn/lNo9kocvGszywnIeX7LJ6jjKSVr0bbAgZzfhwQGcM7Cr1VGU8phL0xK5cnQSL35XqCdT+Qgt+pNUU9/AovV7OG9QPKFBerSNal/uvzBFT6byIVr0J+nbTaUcqq7XYRvVLunJVL5Fi/4kLczZTUx4MKf20WvDqvZJT6byHVr0J+FQdR1fbihh0pBuBAboJlTtl55M5Ru0pU7CktwSauobmazDNkr9fDLVQ5/ksWr7PqvjqBZo0Z+EBWt2kxgdxoikKKujKGW5oydTdY8K4+b/rKb0kJ5M5W206Fup9FANPxSUkTGsOyJ6yUCl4P9OpjpQVcetb6/Wk6m8jBZ9Ky1at4eGRqNH2yjVzNGTqX4s3KcnU3kZLfpWWpBTxID4CPrH6UWTlWru0rRErtKTqbyOFn0r7CivYvWOA/ppXqnjuO/CFIb2iOKud9ewRU+m8gpa9K3w8drdAFw4VGeqVOpYQgIDeP6qEQQH2rj+9Sz2Ha61OlK7p0XvJGMMH2UXMTI5msToDlbHUcqrJUSF8dLVaRQdOMLv3siiuq7B6kjtmha9kzYWHyJ/b6UeO6+Uk9J6dubfVwxj9Y79/PGdHBob9cxZq2jRO2lBzm4CbcKvB+uwjVLOOn9wN/58wUA+W1/MI4s2WB2n3Qq0OoAvaHRcYOT0frF0Dg+2Oo5SPuW603qxa/8RXv5+KwnRYfz21F5WR2p39BO9E1bt2E/RgSN6tI1SJ0FE+MukFCakxPHgJ3ksztXDLj1Ni94JC3KKCA2yMT4lzuooSvmkAJvw76nDGZoYxR/ezmb1Dr3AuCdp0Z9AXUMjn67dw/iUeMJDdKRLqZMVFhzAnBnpxEeGcv3rWWwvP2x1pHZDi/4Evs8vY39VHRlDu1sdRSmfF9MxhFevGYkxhmteXanH2HuIU0UvIhNFZJOIFIjIPS0sDxGRdxzLV4hIsuPx8SKySkTWOb6e7dr47rcgp4jIsCDO6N/F6ihK+YXeXTry8ox0Pcbeg05Y9CISADwHnA+kANNEJKXZatcB+40xfYEngUcdj5cBFxpjBgMzgLmuCu4JVbX1LMkr4YLB3QgO1P/8KOUqeoy9ZznTXqOAAmNMoTGmFpgHZDRbJwN43XH7PeAcERFjTLYxZrfj8VwgTERCXBHcE77csJeq2gYyhumwjVKupsfYe44zexcTgJ1N7u8CRh9rHWNMvYhUADHYP9EfdQmw2hjzi6sSiMhMYCZAUlKS0+HdbWFOEfGdQhmV3NnqKEr5JT3G3jM8Mh4hIqnYh3NuaGm5MWa2MSbdGJPepYt3jIXvP1zLt5tKmTysOzabXmBEKXfQY+w9w5miLwJ6NLmf6HisxXVEJBCIBMod9xOBD4GrjTFb2hrYUz5bX0x9o2GyHm2jlFvpMfbu50zRrwT6iUgvEQkGpgILm62zEPvOVoBLga+NMUZEooBPgXuMMT+4KrQnLMgpok+XcFK7d7I6ilJ+T4+xd68TFr0xph64BVgMbADmG2NyReRBEZnsWG0OECMiBcAdwNFDMG8B+gL3iUiO409Xl/8rXGz3gSP8tG0fGcMS9LqwSnmIHmPvPmKMdx3WlJ6ebrKysizNMDtzC48s2si3d51Fcmy4pVmUam9Wbd/HtJdWMDA+gtevHUVUB51I0BkissoYk97SMj04vAULcnYztEeUlrxSFkjr2ZlZV41gw55DTJ39I6WHfnGgnmolLfpmCvYeInf3QZ3yQCkLnTMwjleuGcn28iouf3E5RQeOWB3Jp2nRN7MwZzc2gUlD9AIjSlnptH6xvHn9KMoqa7j8heVsK9MdtCdLi74JYwwL1uxmXJ9YunYKtTqOUu1eWs/OvP27MRypa+CyF5ezqfiQ1ZF8khZ9E2t2VbC9vIrJOuWBUl5jUEIk828Yg03gitnLWbvrgNOk8ssAAAxNSURBVNWRfI4WfRMLcooIDrQxcVC81VGUUk307RrBuzeMIyI0kCtfWsGKwnKrI/kULXqHhkbDx2v2cPYpXekUGmR1HKVUM0kxHXj3hnHEdQphxqs/8e2mvVZH8hla9A7Lt5RTVlmjM1Uq5cXiI0OZf8NYesd25HdvZPHZuj1WR/IJWvQOC3KKiAgJ5FcDvP7EXaXatZiOIbw9cwxDEqP4/VureX/VLqsjeT0temBzySEWrdvDhNR4QoMCrI6jlDqByLAg5l43irF9Yrjz3TXMXb7N6kherd0X/Y7yKqa/vILwkEBuP7ef1XGUUk7qEBzInBkjOXdgHH9ZkMusb31mclyPa9dFv/dgNdPnrKC2oZE3rx9Nj84drI6klGqF0KAAZk0fweSh3Xn08408vngj3jZ/lzdw5gpTfmn/4Vqmz1lBeWUN//ndGPrHRVgdSSl1EoICbDx5xTDCQwJ47pstHK5p4L5JKXrBoCbaZdFX1tRzzWsr2VZexWu/HcmwHlFWR1JKtUGATXjkosGEBwfy8vdbOVRdz8MXDdJ9bg7truir6xqY+UYW64sqeGF6GuP6xFodSSnlAiLCn389kIjQIJ78cjPrig7w76nDGdhNLx7Ursbo6xsaufXtbJZtKeeflw1hfEqc1ZGUUi4kItx2bj9e++1I9h2uI+PZH5jz/VYaG9v3uH27KfrGRsPd763li7wSHsxI5aLhiVZHUkq5yVmndGXx7adzRv9YHvokjxmv/kTJwWqrY1mmXRS9MYYHP8njg+wi7hzfn6vHJlsdSSnlZjEdQ3jp6nQevmgQK7ftY+JTmSzOLbY6liXaRdE/+cVmXlu2jd+d3otbzu5rdRyllIeICFeN7sknt55OQnQYN8xdxb0frKWqtt7qaB7l90X/8tJCnv66gCvSe/CnCwbqxb6Vaof6du3IBzedyo1n9mHeyp1Mevr7djXdsV8X/fyVO/nbpxu4YHA8j1w8WEteqXYsONDGPecP4K3r7Rcyufj5ZTz3TQEN7WBHrd8W/aJ1e7jng7Wc3i+WJ68YRoCePKGUAsb2ieHz287gvEHxPL54E9Ne+tHvr0nrl0WfubmU2+ZlMzwpmhd/k0ZIoJ40oZT6P5Edgnh22nCeuGwouUUVTHwqk4Vrdlsdy238ruhXbd/HDXNX0bdrBK9cM5IOwe3unDCllBNEhEvSEll02+n07dqRP7ydzR3v5HCous7qaC7nV0Wft/sg17y6kvjIUN64dhSRYXqlKKXU8fWMCefdG8Zy+7n9+CiniIlPLeXtn3ZQXddgdTSX8Zui31p2mKtfWUHHkEDmXjeKLhEhVkdSSvmIwAAbt5/bn3dvHEdUhyDu/WAdY//+Ff9cvMkvTrTym3GNQJuQ1LkDj106lMRonW5YKdV6aT2j+eTW0/hp6z7mfL+V574t4MXMLUwa0p1rT+3F4MRIqyOeFPG2uZvT09NNVlbWST3XGKOHUCqlXGZHeRWvLtvK/JU7OVzbwKjkzlx7WjLjU+K97kg+EVlljElvcZk/Fb1SSrnDweo65q/cyWvLtrFr/xESo8O4Zlwyl4/sQadQ79gXqEWvlFIu0NBo+CKvmFe+38ZP2/bRMSSQy9ITuWZcMj1jwi3NpkWvlFIutm5XBa/8sJWP1+ymwRjOHRjHtaf2YlSvzpYM67S56EVkIvBvIAB42Rjzj2bLQ4A3gDSgHLjCGLPNsexe4DqgAfiDMWbx8V5Li14p5UtKDlYzd/l2/rNiO/ur6ggOsJEU04FeseH0jg2nd5dwesV2pFdsOLEdg922H7FNRS8iAcBmYDywC1gJTDPG5DVZ52ZgiDHmRhGZClxkjLlCRFKAt4FRQHfgS6C/MeaYB6hq0SulfFF1XQOfry9mQ/FBtpYeZmvZYbaXV1Hb0PjzOhEhgfTqEk6v2P/70zu2I8mxHYho41j/8YremcMrRwEFxphCx182D8gA8pqskwE84Lj9HvCs2H9tZQDzjDE1wFYRKXD8fctP5h+ilFLeKjQogCnDE5hCws+PNTQadh84QmHZYbaWVrK17DCFZYdZtX0/C9fspunn7C4RIWQM7c7/TkpxeTZnij4B2Nnk/i5g9LHWMcbUi0gFEON4/Mdmz01o9lxEZCYwEyApKcnZ7Eop5dUCbEKPzh3o0bkDZ/bv8l/Lqusa2LGvikLHp/+tZZV0iwpzSw6vOGHKGDMbmA32oRuL4yillNuFBgXQPy6C/nERbn8tZ6ZAKAJ6NLmf6HisxXVEJBCIxL5T1pnnKqWUciNnin4l0E9EeolIMDAVWNhsnYXADMftS4GvjX0v70JgqoiEiEgvoB/wk2uiK6WUcsYJh24cY+63AIuxH175ijEmV0QeBLKMMQuBOcBcx87Wfdh/GeBYbz72Hbf1wO+Pd8SNUkop19MTppRSyg8c7/BKv5mmWCmlVMu06JVSys9p0SullJ/ToldKKT/ndTtjRaQU2N6GvyIWKHNRHFfSXK2juVpHc7WOP+bqaYzp0tICryv6thKRrGPtebaS5modzdU6mqt12lsuHbpRSik/p0WvlFJ+zh+LfrbVAY5Bc7WO5modzdU67SqX343RK6WU+m/++IleKaVUE1r0Sinl53yy6EVkoohsEpECEbmnheUhIvKOY/kKEUn2QKYeIvKNiOSJSK6I3NbCOmeJSIWI5Dj+3OfuXE1ee5uIrHO87i9mjRO7px3bbK2IjHBznlOabIccETkoIrc3W8dj20tEXhGRvSKyvsljnUXkCxHJd3yNPsZzZzjWyReRGS2t4+Jcj4vIRsf36UMRiTrGc4/7PXdDrgdEpKjJ9+uCYzz3uD+/bsj1TpNM20Qk5xjPdef2arEfPPYeM8b41B/sUyVvAXoDwcAaIKXZOjcDLzhuTwXe8UCubsAIx+0I7BdUb57rLOATi7bbNiD2OMsvAD4DBBgDrPDw97QY+wkflmwv4AxgBLC+yWOPAfc4bt8DPNrC8zoDhY6v0Y7b0W7ONQEIdNx+tKVcznzP3ZDrAeAuJ77Xx/35dXWuZsufAO6zYHu12A+eeo/54if6ny9WboypBY5erLypDOB1x+33gHNERNwZyhizxxiz2nH7ELCBFq6P68UygDeM3Y9AlIh089BrnwNsMca05YzoNjHGZGK/lkJTTd9HrwNTWnjqecAXxph9xpj9wBfARHfmMsYsMcbUO+7+iP3KbR51jO3lDGd+ft2Sy9EBlwNvu+r1nHWcfvDIe8wXi76li5U3L9T/ulg5cPRi5R7hGCoaDqxoYfFYEVkjIp+JSKqnMgEGWCIiq8R+MfbmnNmu7jKVY//wWbW9AOKMMXsct4uBuBbWsXK7AVyL/X9iLTnR99wdbnEMKb1yjGEIK7fX6UCJMSb/GMs9sr2a9YNH3mO+WPReTUQ6Au8DtxtjDjZbvBr78MRQ4BngIw9GO80YMwI4H/i9iJzhwdc+JrFfnnIy8G4Li63cXv/F2P8P7VXHIovIn7Ffue0/x1jF09/zWUAfYBiwB/swiTeZxvE/zbt9ex2vH9z5HvPFom/LxcrdSkSCsH8T/2OM+aD5cmPMQWNMpeP2IiBIRGLdncvxekWOr3uBD7H/F7opqy7kfj6w2hhT0nyBldvLoeTo8JXj694W1rFku4nINcAk4CpHQfyCE99zlzLGlBhjGowxjcBLx3g9q7ZXIHAx8M6x1nH39jpGP3jkPeaLRd+Wi5W7jWP8bw6wwRjzr2OsE390X4GIjMK+/T3xCyhcRCKO3sa+M299s9UWAleL3Rigosl/Kd3pmJ+yrNpeTTR9H80AFrSwzmJggohEO4YqJjgecxsRmQjcDUw2xlQdYx1nvueuztV0n85Fx3g9Z35+3eFcYKMxZldLC929vY7TD555j7ljD7O7/2A/QmQz9r33f3Y89iD2Nz5AKPahgALgJ6C3BzKdhv2/XWuBHMefC4AbgRsd69wC5GI/0uBHYJyHtldvx2uucbz+0W3WNJsAzzm26Tog3QO5wrEXd2STxyzZXth/2ewB6rCPgV6Hfb/OV0A+8CXQ2bFuOvByk+de63ivFQC/9UCuAuxjtkffZ0ePMOsOLDre99zNueY63jtrsRdYt+a5HPd/8fPrzlyOx187+r5qsq4nt9ex+sEj7zGdAkEppfycLw7dKKWUagUteqWU8nNa9Eop5ee06JVSys9p0SullJ/ToldKKT+nRa+UUn7u/wPWygRrs0vJwwAAAABJRU5ErkJggg==\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KTYeeBp4LDuy"
      },
      "source": [
        "# %reload_ext tensorboard\n",
        "# %tensorboard --logdir lightning_logs/\n"
      ],
      "execution_count": 102,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### VERIFY THAT EMA MODEL SAVING WORKS AS INTENDED\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "lxmH-Zb7jfAK"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DTvtvxzPCVcD",
        "outputId": "5cd99336-261c-4ae3-f259-b41bed64a4b0"
      },
      "source": [
        "# original weights before ema\n",
        "toxic_comment_model.pytorch_model.classifier.weight\n",
        "\n"
      ],
      "execution_count": 103,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Parameter containing:\n",
              "tensor([[ 0.0316, -0.0078,  0.0104,  ..., -0.0257,  0.0246, -0.0016],\n",
              "        [ 0.0284,  0.0080, -0.0036,  ..., -0.0382, -0.0260, -0.0281],\n",
              "        [ 0.0157,  0.0159,  0.0282,  ..., -0.0107, -0.0175,  0.0080],\n",
              "        [ 0.0080, -0.0246,  0.0032,  ..., -0.0113,  0.0063, -0.0386],\n",
              "        [-0.0117,  0.0171, -0.0172,  ..., -0.0439,  0.0406,  0.0099],\n",
              "        [-0.0527, -0.0562,  0.0380,  ..., -0.0491, -0.0010, -0.0339]],\n",
              "       requires_grad=True)"
            ]
          },
          "metadata": {},
          "execution_count": 103
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vbEYAlSHCgsd",
        "outputId": "65706e55-0cef-4665-eb56-61244a825b87"
      },
      "source": [
        "bert_model2 = BertModel(len(LABEL_COLUMNS))\n",
        "toxic_comment_model.ema.copy_to(bert_model2.parameters())\n",
        "print(\"after ema weights\")\n",
        "print(bert_model2.classifier.weight)"
      ],
      "execution_count": 104,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "after ema weights\n",
            "Parameter containing:\n",
            "tensor([[ 0.0316, -0.0079,  0.0104,  ..., -0.0257,  0.0246, -0.0016],\n",
            "        [ 0.0257,  0.0053, -0.0009,  ..., -0.0409, -0.0233, -0.0308],\n",
            "        [ 0.0122,  0.0125,  0.0317,  ..., -0.0142, -0.0140,  0.0045],\n",
            "        [ 0.0056, -0.0270,  0.0056,  ..., -0.0137,  0.0087, -0.0411],\n",
            "        [-0.0115,  0.0173, -0.0174,  ..., -0.0437,  0.0404,  0.0100],\n",
            "        [-0.0605, -0.0640,  0.0457,  ..., -0.0569,  0.0067, -0.0417]],\n",
            "       requires_grad=True)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# verify model loading\n",
        "bert_model3 = BertModel(len(LABEL_COLUMNS))\n",
        "bert_model3.load_state_dict(torch.load(\"test.pth\"))\n",
        "print(bert_model3.classifier.weight)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 164
        },
        "id": "12PiDGipjvmw",
        "outputId": "0c5c7de0-1e7c-4e8d-9398-cd13fb8243e1"
      },
      "execution_count": 105,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Parameter containing:\n",
            "tensor([[ 0.0316, -0.0079,  0.0104,  ..., -0.0257,  0.0246, -0.0016],\n",
            "        [ 0.0257,  0.0053, -0.0009,  ..., -0.0409, -0.0233, -0.0308],\n",
            "        [ 0.0122,  0.0125,  0.0317,  ..., -0.0142, -0.0140,  0.0045],\n",
            "        [ 0.0056, -0.0270,  0.0056,  ..., -0.0137,  0.0087, -0.0411],\n",
            "        [-0.0115,  0.0173, -0.0174,  ..., -0.0437,  0.0404,  0.0100],\n",
            "        [-0.0605, -0.0640,  0.0457,  ..., -0.0569,  0.0067, -0.0417]],\n",
            "       requires_grad=True)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "conclusion: model saver with EMA works as intended and can be loaded directly to pytorch module"
      ],
      "metadata": {
        "id": "P3zM09hPqkQE"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_GEsu9KydC93"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}