lit_train_12_attmask_mean_do.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from forgebox.imports import *\n",
    "from forgebox.ftorch.cuda import CudaHandler\n",
    "from transformers import AutoModel,AutoTokenizer\n",
    "from typing import Dict, List\n",
    "import pytorch_lightning as pl\n",
    "from torch.utils.data.dataloader import DataLoader, default_collate\n",
    "from copy import deepcopy\n",
    "\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T13:46:18.835247Z",
     "iopub.status.busy": "2021-07-22T13:46:18.834854Z",
     "iopub.status.idle": "2021-07-22T13:46:18.839667Z",
     "shell.execute_reply": "2021-07-22T13:46:18.838807Z",
     "shell.execute_reply.started": "2021-07-22T13:46:18.835213Z"
    }
   },
   "outputs": [],
   "source": [
    "tag = \"roberta-large\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained = \"weights/pre_rbtlg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T13:46:18.994272Z",
     "iopub.status.busy": "2021-07-22T13:46:18.993895Z",
     "iopub.status.idle": "2021-07-22T13:46:42.849012Z",
     "shell.execute_reply": "2021-07-22T13:46:42.847886Z",
     "shell.execute_reply.started": "2021-07-22T13:46:18.994243Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at weights/pre_rbtlg were not used when initializing RobertaModel: ['lm_head.layer_norm.weight', 'lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.layer_norm.bias', 'lm_head.bias', 'lm_head.decoder.weight', 'lm_head.decoder.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",
      "Some weights of RobertaModel were not initialized from the model checkpoint at weights/pre_rbtlg and are newly initialized: ['roberta.pooler.dense.weight', 'roberta.pooler.dense.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('./rbt_lg_pre')\n",
    "model = AutoModel.from_pretrained(pretrained)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from forgebox.config import Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG = Config(\n",
    "    learning_rate=1e-6,\n",
    "    top_lr_expansion=1e3,\n",
    "    max_epochs = 15,\n",
    "    frozen_years = 3,\n",
    "    drop_out_ratio=.2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A look at the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:19:15.307014Z",
     "iopub.status.busy": "2021-07-22T14:19:15.306588Z",
     "iopub.status.idle": "2021-07-22T14:19:15.374632Z",
     "shell.execute_reply": "2021-07-22T14:19:15.373623Z",
     "shell.execute_reply.started": "2021-07-22T14:19:15.306980Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>url_legal</th>\n",
       "      <th>license</th>\n",
       "      <th>excerpt</th>\n",
       "      <th>target</th>\n",
       "      <th>standard_error</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>c12129c31</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>When the young people returned to the ballroom...</td>\n",
       "      <td>-0.340259</td>\n",
       "      <td>0.464009</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>85aa80a4c</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>All through dinner time, Mrs. Fayre was somewh...</td>\n",
       "      <td>-0.315372</td>\n",
       "      <td>0.480805</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>b69ac6792</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>As Roger had predicted, the snow departed as q...</td>\n",
       "      <td>-0.580118</td>\n",
       "      <td>0.476676</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>dd1000b26</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>And outside before the palace a great garden w...</td>\n",
       "      <td>-1.054013</td>\n",
       "      <td>0.450007</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>37c1b32fb</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Once upon a time there were Three Bears who li...</td>\n",
       "      <td>0.247197</td>\n",
       "      <td>0.510845</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          id url_legal license  \\\n",
       "0  c12129c31       NaN     NaN   \n",
       "1  85aa80a4c       NaN     NaN   \n",
       "2  b69ac6792       NaN     NaN   \n",
       "3  dd1000b26       NaN     NaN   \n",
       "4  37c1b32fb       NaN     NaN   \n",
       "\n",
       "                                             excerpt    target  standard_error  \n",
       "0  When the young people returned to the ballroom... -0.340259        0.464009  \n",
       "1  All through dinner time, Mrs. Fayre was somewh... -0.315372        0.480805  \n",
       "2  As Roger had predicted, the snow departed as q... -0.580118        0.476676  \n",
       "3  And outside before the palace a great garden w... -1.054013        0.450007  \n",
       "4  Once upon a time there were Three Bears who li...  0.247197        0.510845  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_df = pd.read_csv(\"train.csv\")\n",
    "test_df = pd.read_csv(\"train.csv\").head(20)\n",
    "display(train_df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch = tokenizer(list(train_df.head()[\"excerpt\"]), return_tensors=\"pt\", padding=True)\n",
    "\n",
    "# token_ids = batch[\"input_ids\"]\n",
    "# attention_mask = batch[\"attention_mask\"]\n",
    "\n",
    "# attention_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T13:51:57.676060Z",
     "iopub.status.busy": "2021-07-22T13:51:57.675670Z",
     "iopub.status.idle": "2021-07-22T13:51:57.840797Z",
     "shell.execute_reply": "2021-07-22T13:51:57.839727Z",
     "shell.execute_reply.started": "2021-07-22T13:51:57.676030Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/anaconda3/lib/python3.7/site-packages/matplotlib/axes/_base.py:405: FutureWarning: Support for multi-dimensional indexing (e.g. `obj[:, None]`) is deprecated and will be removed in a future version.  Convert to a numpy array before indexing instead.\n",
      "  x = x[:, np.newaxis]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7d6fae9990>]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXMAAAD2CAYAAAAksGdNAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAAkGUlEQVR4nO3deXiV1aHv8e/aGUkgIRNkYAgQxgAhGJEhTqiIMhRqbwdrj0O99vScDvfY2tae9jntc9vT6u1pb08n9dRTh2ttawVUVBRRkRkT5kFGk5CQhMwDmZN1/8iGFghTspP33Xv/Ps+Th513h+xf34afb9Ze613GWouIiPg3j9MBRESk71TmIiIBQGUuIhIAVOYiIgFAZS4iEgBCnXjRxMREm56e7sRLi4j4rfz8/EprbVJPzzlS5unp6eTl5Tnx0iIifssYU3ix5zTMIiISAFTmIiIBQGUuIhIAVOYiIgFAZS4iEgBU5iIiAUBlLiISAFTm4je6uiwvbi+iua3T6SgirqMyF7+x4Wglj67Yy0v5J5yOIuI6KnPxG5uOVgKw8Uilw0lE3EdlLn7jTIlvOV5FR2eXw2lE3EVlLn6hsrGVA6X1TEmJoaGlgz0ldU5HEnEVlbn4hc3HqgB4ZOFEADZpqEXkHCpz8Qsbj1QQOyiMG8YnkZkaw4ajKnORv6cyF9ez1rLxSCVzxyUQ4jHkjk9kZ1ENp1s7nI4m4hoqc3G9jytPc7KuhXkZiQDkZiTS3mnZ/nG1w8lE3ENlLq53Zkri9eO7y/za9HjCQz1s1FCLyFkqc3G9DUcqGRE3iFHxUQBEhoVwbXqc5puL/B2VubhaR2cXW45Xcf34RIwxZ4/nZiRxqLyBUw0tDqYTcQ+VubjanpI6Glo6zo6Xn5Hr/Xzz0SonYom4jspcXG3TkUqMgbnjzi3zzNQYhkaFsUFDLSKAylxcbsPRSjJTY4iPDj/nuMdjmDcukU1HK7HWOpROxD1U5uJap1s72FlUQ25GUo/P545PpKy+hWMVjQOcTMR9VObiWts/rqa9054dHz/fmeOa1SKiMhcX23i0kvBQDznpcT0+PzI+ilHxUZpvLoLKXFxs45FKZqXHExkWctGvyR2fyNbj1bTrlrgS5FTm4kqnGlo4VN5wwZTE8+VmJNLY2sHuE7UDE0zEpVTm4krnL+G/mLnjEjAGDbVI0FOZiyttPFJFXFQYU1JiLvl1Q6PCmZYWe7b8RYKVylxcx1rLxqMVzM1IxOMxl/363IxEdhbV0qhb4koQU5mL6xyraKS8vvWiUxLPl5uRSEeXZdtxLe2X4KUyF9c5s0T/Sst85ug4IsM8WtovQe2yZW66PWuM2WqMedUYM9gYs9oYs9sY87z3+cjzjw1EeAlMm45WMjohipHeW95eTvctceM1bi5B7UquzOcBodba2UAM8ABQbK3NAuKA24B7ejgmctXaO7vYerz6iq/Kz7h+fCJHTjVSVqdb4kpwupIyLwd+6X3cBvwAWOv9/F3gZmB+D8dELtDY2nHJG2PtPtH9RubVlvmZ+ei6Opdgddkyt9YesdZuN8YsB8KBfKDO+3Q9EA8k9HDsHMaYh4wxecaYvIqKCp+EF/9y4GQ9WT98m5t/9j4/efMgO4pq6Oo6t9g3XOSWt5czOTmGhOhwzTeXoBV6JV9kjFkKfB1YAjwBxHqfigUqgcE9HDuHtfYp4CmAnJwc3bM0CL2Uf4IQYxgZH8XTGz7myfXHGR4Twe2ZySzMTGbWmO5x7+lpscRGhV3V9/Z4DHMzEtnovSWu3raRYHPZMjfGJAOPAAuttaeNMeuABcDLdA+v/AIY1cMxkbM6Ort4bfdJbpk8jN/dcw11ze28+1E5a/aV8Ze8Ezy3pZC4qDDqWzr4xxvH9uo1rs9I5LXdJzlc3sjE5CE+/l8g4m5XcmV+L5ACvOW92nkeSDPG7AF2A+voHn755HnHRM7aeLSSysY2lmWnARA7KIzl2SNYnj2CprYOPjhcwZp9ZeQX1bBoWmqvXuOGCUl4DKzYWcyjd0z2ZXwR17tsmVtrHwMeO+/wk+d93gos9lUoCTyrdpYQOyiMmyZeuNFEVHgoC6emsHBqSp9eIzk2kjumpfDHbUV8bf54oiOuaBRRJCBo0ZD0u9OtHby1v5xF01OICL347Wx94cHcMTS0dPCXvBP9+joibqMyl3739oEymts7We4dYulP2aPiuGZ0HP+96WM6u/Q+uwQPlbn0uxU7ShgRN4hrRvW8Y5CvPZg7hhPVzaw9UDYgryfiBipz6Ven6lvYdLSS5dlpV3QHRF9YkJnMyPhB/H7DxwPyeiJuoDKXfvXq7pN0WfjEjP4fYjkjxGO4f+4Y8gpr2FlUM2CvK+Iklbn0q1W7Spg+IpaMYYMH9HU/fe1IhkSG8vRGXZ1LcFCZS785Ut7AvpJ6lg3gVfkZgyNCuXvWKN7cV0ZxTdOAv77IQFOZS79ZtauEEI9hSVbvFgH11b1z0wF4dnOBI68vMpBU5tIvurosq3aeJDcjkaQhEY5kSB06iEXTUvjT9hM0tLQ7kkFkoKjMpV98WFBNSW3zgMwtv5QHrx9DQ2sHf/5Qi4gksKnMpV+s2lVCVHgICzKHO5pj+oihzEqP5w+bCujo7HI0i0h/UpmLz7W0d7J6TykLM5OJCnf+/ihfvH4MJbXNvLW/3OkoIv1GZS4+9/6hUzS0dJy9Q6LTbp08nNEJUfx+43Gno4j0G5W5+NzKnSUkDYlg7rgEp6MA3YuIvpg7hp1FteQXahGRBCaVufhUbVMb7350iqVZqYSGuOfH61PXjCB2UBhP6+pcApR7/rVJQHh9byntndbxWSzniwoP5e7rRrFmXxn7Suou/xdE/IzKXHxq1c4SMoYNJjM1xukoF7hvbjrx0eEs/+0mfr72MK0dnU5HEvEZ56caiF9p7ejkVH0rpXUtlNW3UF7XQmldC+X13Z/nF9bwyO0TXbmh8vCYSN76Xzfwo9cP8p/rjvDG3lJ++slp5KTHOx1NpM+MtQN/A/+cnBybl5c34K8rvdfa0clv3zvG79Yfo63j3PnaUeEhJMdEkhwbyYi4QXx74SQSBjuz6vNKvX/oFP+6ch8ltc18YfZovrVwIkMiw5yOJXJJxph8a21Oj8+pzOVy8gur+fbLezl6qpHF01O4YULS2fJOjo1kSESoK6/EL+d0awc/e/sQz2wuIDkmkh8tm8otk51d5CRyKSpz6ZWGlnb+z1uHeH5rIamxg/jR8qncPHGY07F8bmdRDd95eS+HyhtYPD2Fx+6ars2gxZUuVeb6iZUerTtYzvdW7aOsvoX75qbzzQUTA7bgskfF8dpXc3li/TF+vvYwU1Jj+KebMpyOJXJVNJtFzlHR0MpX/riDLz6bR0xkGCu+PJd/W5IZsEV+Rnioh6/dMp6c0XGs3FGCE7+xivSFylzOKqtrYcEv1vP2/nIevm0Cr301l+wB2oTZLZZlp3HkVCMHSuudjiJyVVTmctZ7h05R09TOiw/N5mu3jCc8NPh+PBZNSyEsxLBqZ4nTUUSuSvD9a5WLyi+sIT46nJmjhjodxTFx0eHcNHEYr+w6SWeXhlrEf6jM5awdhTXMHBXnl9MMfWl5dhqnGlrZfKzS6SgiV0xlLgBUNbZyvPI014wOrjHynsyfNIwhEaGs1FCL+BGVuQCwo6gWgJx0lXlkWAh3TkvhrX1lNLV1OB1H5IqozAWAvMJqwkIM09JinY7iCstnpnG6rZO1B7Q7kfgHlbkA3ePlU9NiiQwLcTqKK8xKjyc1NlKzWsRvqMyF1o5OdhfXcU2QzSm/FI/H8InsND44UkllY6vTcUQuS2Uu7D9ZT1tHl8bLz7M8O43OLsvq3SedjiJyWSpzIb+ge1/MmZrJco4Jw4cwJSVGs1rEL6jMhfzCGkbFRzFsSKTTUVxneXYau4vrOFbR6HQUkUu6ojI3xoQZY17zPl5ojCk2xmz0fkw0xkQaY1YbY3YbY543wb7qxI9Ya8krrNH88otYOiMVj4FXdHUuLnfZMjfGDALygdv+7vDvrLW53o9DwD1AsbU2C4g772vFxU5UN1PZ2KohlosYHhPJvIxEVu7SnRTF3S5b5tbaZmvtdKD47w7fZYzZbox52XsVPh9Y633uXeBm30eV/pBfVA1Ajsr8opbNSONEdTM7imqcjiJyUb0ZMz8GfN9aOwtIAW4EEoA67/P1wAU75BpjHjLG5Blj8ioqKnqbV3wsr6CGIRGhTBg+xOkornX71GQiwzx6I1RcrTdlXg28431cAAwDKoEzSwdjvZ+fw1r7lLU2x1qbk5SU1IuXlf6QX1jDjFFDCfHobY6LGRwRyoIpyazeU3rBZtYibtGbMn8Y+KwxxgNMBfYB64AF3ufnA+/5Jp70p/qWdg6VN+jNzyuwPDuN2qZ23j90yukoIj3qTZn/Grgf2AastNYeAF4A0owxe+i+cl/nu4jSX3YV1WIt5Iy+YFRMzpM7PpGE6HBW7dJQi7jTFW/saK3N8P5ZCtx03nOtwGKfJpN+l1dYg8fAjCDejOJKhYV4WJKVyh+3F1HX3E7soDCnI4mcQ4uGgtiOwhomJccwOMA3a/aV5dlptHV0sWJH8eW/WGSAqcyDVEdnFzuLtFjoakwfEcus9HieXH+c1o5Op+OInENlHqQOlTdwuq1TZX4VjDF87ZbxlNW38FKers7FXVTmQSq/sHsBjMr86szLSGDmqKH87v1jmqYorqIyD1L5hTUMGxLBiLhBTkfxK2euzktqmzV2Lq6iMg9S+YU15KTHoXuiXb0bJySRNSKW37x/lPZOXZ2LO6jMg1B5fQvFNc3M1M5CvXLm6vxEdbO2lRPXUJkHoTPj5TnpWizUW/MnDSMzNYbfvHeUDl2diwuozINQXkENEaEepqTEOB3Fb525Oi+oauJVbSsnLqAyD0L5RTVkjRxKeKj+7++L2yYPZ1LyEH797lE6u3Svc3GW/jUHkILK0yz4xXp+v+H4Rculua2T/SV1mpLoAx6P4eu3jOd45WlW79HVuThLZR5AXso/weHyRn70+kE++dtNHCytv+Br9hTX0tFluUZvfvrE7ZnJTBg+mF+9e5QuXZ2Lg1TmAcJayxt7y5iXkcAvPzuDEzXNLPnVRn721iFa2v+29DzP++antonzDY/H8NX54zl6qpE395U5HUeCmMo8QBworefjytMsmpbKJ2ak8c7DN7J0Riq/fu8od/5yA9uOVwHdN9camxRNfHS4w4kDx53TUhiXFM2v3j2iq3NxjMo8QLyxt5QQj+H2zOEAxEeH8/NPz+C5B2bR1tnFZ57ayndX7iW/qEb7ffpYiPfq/KOyBt4+UO50HAlSKvMAYK3l9T2lzBmbQMLgiHOeu2FCEm//yw08mDuGP20vorapXW9+9oPF01MYkxjNf647grW6OpeBpzIPAAdK6ymoamLR9JQen48KD+V7i6ew8p/m8fnrRrEws+evk94LDfHwzzdncKC0nvcPacNyGXgq8wDw+p4zQyzJl/y6rJFD+fHyacRGaZec/vCJGanER4ezQkv8xQEqcz/XPYullLnjEvSmpsPCQjzcMTWZdw6U09TW4XQcCTIqcz+3/6R3iGWahk7cYElWKs3tnazVG6EywFTmfu517yyWBZcZYpGBMSs9nuSYSF7bXep0FAkyKnM/piEW9/F4DIunp7D+8CnqmtqdjiNBRGXux/afrKdQQyyus3RGKu2dljX7dXUuA0dl7sde33tls1hkYE1Li2V0QpRujSsDSmXup84sFJqXkUichlhcxRjD0qxUthyr4lRDi9NxJEiozP3U/pP1FFU3sWiarsrdaGlWKl0W3tijoRYZGCpzP7V6TymhHsOCKSpzNxo/fAiTkodoqEUGjMrcD1lreX3vSeZqiMXVlmSlsqOolhPVTU5HkSCgMvdD+0rqOVHdzGLNYnG1pVmpQPdvUSL9TWXuh1bvPdk9xOK93a2408j4KLJHDdVQiwwIlbmfObNQaF5GIkOjNMTidkuzUjlYWs/RUw1OR5EApzL3M3tL6jhR3ayFQn5i0fQUPAZe3aWrc+lfKnM/8/reUg2x+JFhQyKZPTaB1/aUatMK6Vcqcz9yZqFQ7ngNsfiTpVmpfFx5mn0l9U5HkQCmMvcjW45XUVzTzJ0aYvErd0xNISzE8OpubVoh/eeKytwYE2aMec37ONIYs9oYs9sY87zpdsGx/o0dfJrbOvnuir2MjB/E4otsDyfuFBsVxo0Tkli9p5SuLg21SP+4bJkbYwYB+cBt3kP3AMXW2iwgznu8p2PiQ4+/9REFVU08flcWUeGhTseRq7QkK5XSuhbyCmucjiIB6rJlbq1tttZOB4q9h+YDa72P3wVuvsixcxhjHjLG5Blj8ioqtOHt1dj+cTXPbC7g3jmjmTMuwek40gu3Th5OZJhHQy3Sb3ozZp4A1Hkf1wPxFzl2DmvtU9baHGttTlJSUm+yBqWmtg4e+etuRsZF8e07JjkdR3opOiKUWycP5429ZbR3djkdRwJQb8q8Eoj1Po71ft7TMfGBx9ccorCqicc/NV3DK35uaVYq1afb2HysyukoEoB6U+brgAXex/OB9y5yTPpo2/EqntlcwH1z05k9VsMr/u6GCUmEh3rYcFjDjOJ7vSnzF4A0Y8weoJruIu/pmPRB9/DKHkYnRPGthROdjiM+EBkWwsxRQ9lyXFfm4ntX/Hu7tTbD+2crsPi8p3s6Jn3w2JsfUVTdxJ8fmq3hlQAyZ2wi/3fdYWqb2rTwS3xKi4ZcaMuxKp7dUsh9c9O5TsMrAWXOuASs7Z6hJOJLKnOXOd3awbde3k26hlcCUtbIWCJCPRpqEZ/T7+8u89iajyiuaeYvX5qj4ZUAFBEaQk56HFs0o0V8TFfmLrKvpI7nthRy/9wxXJt+wVR9CRBzxibwUVkD1afbnI4iAURl7iKv7CohLMTw9VvHOx1F+tGZVbzbP9bVufiOytwlrLW8ua+MeRmJxA4KczqO9KNpaUMZFBaioRbxKZW5S+w/WU9xTTN3TE12Oor0s/BQT/e4ud4EFR9SmbvEmn1leEz3DZkk8M0Zl8Dh8kYqG1udjiIBQmXuEm/uK+W6MQkkDI5wOooMgDne9QPbjmu+ufiGytwFjp5q4FjFae6YpiGWYDE1LZbo8BC2HNc96cQ3VOYu8ObeMgAWTFGZB4uwEA/XjonXm6DiMypzF1izv4yZo4aSHBvpdBQZQHPGJnCs4jSn6lucjiIBQGXusKKqJvafrOeOqdrXM9icmW++VfdpER9QmTtszf5SABZqSmLQmZISw5CIUA21iE+ozB22Zl8ZmakxjIyPcjqKDLDQEA+zxsSzVfPNxQdU5g4qq2thR1GtFgoFsTnjEvi48jRldRo3l75RmTvo7QPds1g0xBK8zmwHqKtz6SuVuYPe3FtGxrDBZAwb4nQUccjklBhiIjVuLn2nMndIVWMr2z6uYmGmrsqDWYjHcN3YBN2nRfpMZe6Qdw6W02U1xCLd882LqpsoqW12Oor4MZW5Q97cV8bI+EFkpsY4HUUcdnbcXEMt0gcqcwfUt7Sz6WglCzOTMcY4HUccNil5CHFRYRpqkT5RmTvg3YOnaO+0LNSqTwE8HsN1YxL0Jqj0icrcAW/uK2V4TATZI4c6HUVcYs64BEpqmzlR3eR0FPFTKvMB1tTWwfrDFdyemYzHoyEW6XZm3FxDLdJbKvMBtv5QBS3tXZrFIueYMHwwCdHhehNUek1lPsDW7C8jPjqcWenxTkcRFzHGMNs739xa63Qc8UMq8wHU2tHJuwdPcdvk4YSG6NTLuWaPS6C0roXCKo2by9VTowyg9z6qoKG1g4XaHk56MG9cAsbAPU9v478+OE5dU7vTkcSPqMwH0HNbCkiNjeT6jESno4gLjU0azNP35pA6dBA/fuMgs3+yjkdX7OVQWYPT0cQPhDodIFgcLm9g87EqvrVwooZY5KLmTxrO/EnDOXCynmc3F7BiRzEvbi9iztgE7puXzq2ThxOiWVDSA7XKAHl2cwHhoR4+e+0op6OIH5iSGsNjn5rO1kdv4dsLJ1FYdZovPZ/PDY+/pyt16ZHKfADUNbezYkcJn8hKJT463Ok44kfiosP58k3j+OBbN/PEPTOpb27nvzYcdzqWuJDKfAC8lHeC5vZO7p2b7nQU8VOhIR4WTk3hzmkpvLm3lOa2Tqcjicv0qsyNMQuNMcXGmI3ejyxjzGpjzG5jzPNGd486q6vL8vzWQq4ZHcfUtFin44ifW5adxum2TtYeLHc6irhMX67Mf2etzbXW5gLXAsXW2iwgDrjNJ+kCwPrDFRRWNemqXHziujHxpMZGsmpnidNRxGX6UuZ3GWO2G2NeBm4B1nqPvwvc3OdkAeKZzQUMGxKhTZvFJzwewyey01h/uIKqxlan44iL9LbMjwHft9bOAlKATwJ13ufqgQvWqhtjHjLG5Blj8ioqKnr5sv7leEUj6w9X8PnrRhOm6YjiI8uz0+jssqzeU+p0FHGR3jZMNfCO93EB0AWcGRCOBSrP/wvW2qestTnW2pykpKRevqx/eW5LIWEhhs9dN9LpKBJAJgwfwpSUGFZoqEX+Tm/L/GHgs8YYDzAV+AawwPvcfOA9H2Tza42tHfw1v5hF01IYNiTS6TgSYJZnp7H7RC3HKxqdjiIu0dsy/zVwP7ANWAk8DaQZY/bQfdW+zjfx/NeKHcU0tnbojU/pF0tnpGIMrNp10uko4hK9Ws5vrS0Fbjrv8OI+pwkQ1lqe3VxA1ohYskfFOR1HAtDwmEjmjUtk1c4S/uXW8dpLVrRoqD9sPFrJsYrTuiqXfrU8O42i6iZ2FNU6HUVcQGXeD57dXEBCdDiLpmvDZuk/t09NJjLMw8qdxU5HERdQmftYUVUT6z46xedmjSIiNMTpOBLABkeEsmBKMqv3lNLW0eV0HHGYytzHnt9agMcYPj9bd0eU/rc8O43apnbWHw6OtRtycSpzH2pq6+DPH55gYWYyKbGDnI4jQSB3fCIJ0eFa3i8qc196fM0h6ls6uH9eutNRJEiEhXhYkpXK2oPl1Ldom7lgpjL3kVd2lfDM5gIemDeGnPQL7mYg0m+WZafR1tHFmr1lTkcRB6nMfeBQWQPfeXkv16bH8eidk5yOI0Ema0QsYxOjWaFZLUFNZd5H9S3t/OP/y2dwZCi/uXumbqglA84Yw7LsNLYer6akttnpOOIQNU8fdHVZvvGX3RRVN/Gbu2cyLEb3YBFnLJuRBsCrWt4ftFTmffDEB8dYe6Cc7945mVljNE4uzhmVEMU1o+NYubMYa63TccQBKvNe2nS0kp+9dYjF01N4QLNXxAWWZadxuLyRA6X1TkcRB6jM/05ZXQvfX7WPV3aVXHKa18naZr764k7GJQ3msbum6yZH4gqLp6UQFR7Cf7x9WFfnQahXd00MRNZaHvnrbjYcqeT5rd2bSszLSGRhZjK3TRlOwuAIAFo7OvnyCzto6+jiiS9cQ3SETqG4Q1x0OA/fNoEfvX6QNfvKuGOa7g0UTNREXn/68AQbjlTyw6WZTE2LYc2+MtbsL+M7K/by3ZV7uTY9noVTkzlYWs/uE7U8cc9MxiUNdjq2yDnum5vOih0l/OC1/eSOT2RIZJjTkWSAGCd+HcvJybF5eXkD/roXU1LbzO2/+IBpabG88OB1eDzdwybWWg6U1vOWt9gPl3fv6vKlG8fy6B2TnYwsclG7TtSy/LebuHdOOj9Ymul0HPEhY0y+tTanp+eC/srcWst3Xt5Dl7U8/qnpZ4scuufvZqbGkpkay8MLJnKsopEDJ+u5Y2qyg4lFLm3GyKH8w+zRPLulgOXZaWSNHOp0JBkAQf8G6Ivbu4dXHr1zMiPjoy75teOSBrMkK5VQLQwSl/vG7RMZNiSCR1fspaNTt8cNBkHdSieqm/jx6weYOy6Bz8/SLWslcMREhvGDJZkcKK3nmc0FTseRARC0ZW6t5Tsr9gDw2F3nDq+IBIKFU5OZP2kYP197WMv8g0DQlvkL24rYdLSK7y66/PCKiD8yxvDDpZlYC//2yj7NPQ9wQVnmJ6qb+MkbB8nNSORuDa9IABsZH8W/3Daedw6e4q395U7HkX4UdGXe1WX59st7MMbw07umafWmBLz7541hckoMP3h1P42tHU7HkX4SdGX+wvYiNh+r4rt3TmZEnIZXJPCFhXj49+VTKW9o4T/ePuR0HOknQVXmxysa+ckbB7l+fCKfmzXS6TgiAyZ7VBz3XDeaZzcX8MquEo2fB6CgKfOTtc184entRIaF8FPdHEuC0CMLJzIxOYav/2kXy36zic1HK52OJD4UFGVe1djKPU9vo765necemEXa0EFORxIZcDGRYaz+ai6Pf2o6FQ2t3P37bXzh6W3sK6lzOpr4QMCXeX1LO/f+YTslNc08fd+1TE2LdTqSiGNCPIZP54zk3W/exPcWTWZvSR2Lf7WRr/xxBwWVp52OJ30Q0GXe3NbJg8/k8VFpA0984RrtBiTiFRkWwoPXj+WDb93MV27OYN3BU9z68/X868q91DVd/F7+4l4BW+ZtHV18+YV8Piys5hefmcHNE4c5HUnEdWIiw/jm7RNZ/8hNfG7WKP784Qk+89QWKhtbnY4mVykgy7yzy/LwX3bx/qEK/n35NJZkpTodScTVhsVE8r+XTeWZ+2dRUHWazzy5hbK6FqdjyVUIuDK31vK9VXtZvaeU7945ic9phafIFcsdn8hzD1xHWV0Ln35yC8U1TU5HkisUUGVureWnb37Ei9tP8M83j+OhG8Y5HUnE78waE88L/3M2tU1tfPqJLXpj1E8ETJnvK6nj3j98yJMfHOcLs0fzzQUTnY4k4rdmjBzKiw/NpqWji08/uYUj5Q1OR5LL8PsyL6g8zVdf3MniX21kT3Et31s0mR8uzdSiIJE+ykyN5c8PzQbgM09tZf9JzUd3M5/sAWqMiQT+CowE9gD/YC/xjX2xB+iphhZ+te4oL24vIizEwxdzx/DQjWOJ0Qa2Ij5VUHmaz/9+Gw0t7Tz3xeuYoW3oHDMQe4DeAxRbaxcbY1YDtwFv++h7n6O+pZ2n1h/n6Y0f097ZxWdnjeRr88czLCayP15OJOilJ0bz5y/N5u7/2sY9v9/GomkpJMdGdn/E/O3PoVFh+o3YQb4q8/nAy97H7wI30w9l/t5Hp3j4L7uoaWpnSVYq37htAumJ0b5+GRE5z4i4KF76xzl886XdvHfoFBWNrZz/u3dEqIfhMZFEhPr96G2/+sy1I3nw+rE+/76+KvME4MyAWj1wwbuPxpiHgIcARo3q3XTB9MRoZowcyjcWTNSyfJEBNjwmkue/eB0A7Z1dVDS0UlrXQnl9y9k/y+pa6OjSBtKXkjg4ol++r6/KvBI4066x3s/PYa19CngKusfMe/MiYxKj+cP9s3qbUUR8JCzEQ+rQQaTqpnWu4avfh9YBC7yP5wPv+ej7iojIFfBVmb8ApBlj9gDVdJe7iIgMEJ8Ms1hrW4HFvvheIiJy9fS2s4hIAFCZi4gEAJW5iEgAUJmLiAQAlbmISADwyY22rvpFjakACnv51xPpYVGSi7g9H7g/o/L1jfL1jZvzjbbWJvX0hCNl3hfGmLyL3TXMDdyeD9yfUfn6Rvn6xu35LkbDLCIiAUBlLiISAPyxzJ9yOsBluD0fuD+j8vWN8vWN2/P1yO/GzEVE5EL+eGUuIiLnUZmLiAQAvylzY0ykMWa1MWa3MeZ547LNBo0xC40xxcaYjd6PC3ZbcooxJswY85r3sSvP43kZXXMuTbdnjTFbjTGvGmMGu+n89ZBvsVvOnTdfqDHmJWPMJmPMf7vt56+HfK752btaflPm/G3T6Cwgju5No93md9baXO/HIafDABhjBgH5/O18ue489pAR3HMu5wGh1trZQAzwAO46f+fn68I95w5gGbDbWjsPSAG+grvO3zLOzTcDd52/K+ZPZT4fWOt9fGbTaLe5yxiz3RjzstNXHGdYa5uttdOBYu8h153HHjKCe85lOfBL7+M24Ae46/ydnw/cc+4A1gA/N8aEAkOBmbjr/J2frx53nb8r5k9lfv6m0fEOZunJMeD71tpZdP8X/kaH81yM288juOhcWmuPWGu3G2OWA+F0/wbhmvPXQz7XnDtvvkZrbROwie7/8Ljq56+HfGtx0fm7Gv5U5pfdNNph1cA73scFwDDnolyS288juOxcGmOWAl8HlgCncNn5Oy9fJe46dwnGmAhgLt3DKlNx0fnrId90XHT+roY/lbnbN41+GPisMcZD9w/sPofzXIzbzyO46FwaY5KBR4BF1toGXHb+esjnmnPn9Q3gf1hrO4Em4Me46PxxYb7v4a7zd8X8qczdvmn0r4H7gW3ASmvtAYfzXIzbzyO461zeS/ev228ZYzYCYbjr/J2frwn3nDuA3wAPGGO2AFXA07jr/J2fbzHuOn9XTCtARUQCgD9dmYuIyEWozEVEAoDKXEQkAKjMRUQCgMpcRCQAqMxFRALA/wclvFEBux5ZYgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(train_df.target.groupby(pd.cut(list(train_df.target),\n",
    "                                        bins = list(i*.2-4 for i in  range(40)))).count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:14:57.247786Z",
     "iopub.status.busy": "2021-07-22T14:14:57.247300Z",
     "iopub.status.idle": "2021-07-22T14:14:57.257541Z",
     "shell.execute_reply": "2021-07-22T14:14:57.256295Z",
     "shell.execute_reply.started": "2021-07-22T14:14:57.247749Z"
    }
   },
   "outputs": [],
   "source": [
    "class LitDataset(Dataset):\n",
    "    def __init__(self, df):\n",
    "        self.df = df\n",
    "        self.indices = df.index\n",
    "        \n",
    "    def __len__(self,): return len(self.df)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        row_ = dict(self.df.loc[self.indices[idx]])\n",
    "        row = dict((k,v) for k,v in row_.items() if k in [\"excerpt\",\"id\",\"target\"])\n",
    "        return row\n",
    "    \n",
    "    def collate(self, batch):\n",
    "        batch = default_collate(batch)\n",
    "        batch[\"excerpt\"] = tokenizer(batch[\"excerpt\"], \n",
    "                                     truncation=True,\n",
    "                                     return_tensors='pt', \n",
    "                                     return_token_type_ids=False, \n",
    "                                     padding=True,\n",
    "                                     max_length=model.config.max_position_embeddings\n",
    "                                    )\n",
    "        if 'target' in batch:\n",
    "            batch['target'] = batch['target'].float()\n",
    "        return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:14:38.751413Z",
     "iopub.status.busy": "2021-07-22T14:14:38.750966Z",
     "iopub.status.idle": "2021-07-22T14:14:38.757133Z",
     "shell.execute_reply": "2021-07-22T14:14:38.755698Z",
     "shell.execute_reply.started": "2021-07-22T14:14:38.751378Z"
    }
   },
   "outputs": [],
   "source": [
    "train_ds = LitDataset(train_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:45:51.607175Z",
     "iopub.status.busy": "2021-07-22T14:45:51.606811Z",
     "iopub.status.idle": "2021-07-22T14:45:51.613373Z",
     "shell.execute_reply": "2021-07-22T14:45:51.612144Z",
     "shell.execute_reply.started": "2021-07-22T14:45:51.607145Z"
    }
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 12\n",
    "def create_dl(df, istrain = True):\n",
    "    \"\"\"\n",
    "    Create dataloader from a dataframe\n",
    "    \"\"\"\n",
    "    ds = LitDataset(df)\n",
    "    dl = DataLoader(ds, shuffle=istrain,\n",
    "                    batch_size=BATCH_SIZE if istrain else BATCH_SIZE*2,\n",
    "                    collate_fn = ds.collate)\n",
    "    return dl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:46:35.020204Z",
     "iopub.status.busy": "2021-07-22T14:46:35.019834Z",
     "iopub.status.idle": "2021-07-22T14:46:35.074674Z",
     "shell.execute_reply": "2021-07-22T14:46:35.073908Z",
     "shell.execute_reply.started": "2021-07-22T14:46:35.020174Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': ['035e92ff6',\n",
       "  '61e569be7',\n",
       "  '7a0474321',\n",
       "  '034bfda3f',\n",
       "  '9bd5e9b85',\n",
       "  '590467bd6',\n",
       "  'e21a8d1fb',\n",
       "  '4bff73b96',\n",
       "  '5ac1be4cf',\n",
       "  '9fa6a7f79',\n",
       "  '07cfb6556',\n",
       "  '1e04098af'],\n",
       " 'excerpt': {'input_ids': tensor([[    0,   133, 24696,  ...,     1,     1,     1],\n",
       "         [    0,   250,  2379,  ...,     1,     1,     1],\n",
       "         [    0, 40821,  6721,  ..., 14341,     4,     2],\n",
       "         ...,\n",
       "         [    0,  2847, 24237,  ...,     1,     1,     1],\n",
       "         [    0,  3684,    42,  ...,     1,     1,     1],\n",
       "         [    0,  1711,  3685,  ...,     1,     1,     1]]), 'attention_mask': tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 1, 1, 1],\n",
       "         ...,\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0]])},\n",
       " 'target': tensor([-1.0819,  0.0088,  1.4104, -1.6244, -0.9866, -1.7530, -0.4556, -0.4197,\n",
       "          0.3568, -0.6291, -1.7436, -1.7469])}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a_dl = create_dl(train_df)\n",
    "batch = next(iter(a_dl))\n",
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:37:38.474444Z",
     "iopub.status.busy": "2021-07-22T14:37:38.473793Z",
     "iopub.status.idle": "2021-07-22T14:37:38.481405Z",
     "shell.execute_reply": "2021-07-22T14:37:38.480502Z",
     "shell.execute_reply.started": "2021-07-22T14:37:38.474392Z"
    }
   },
   "outputs": [],
   "source": [
    "def split_k(df, k: int=5):\n",
    "    kmap = np.random.choice(range(k), size=len(df))\n",
    "    return dict((ki, dict(train=(kmap!=ki), valid=(kmap==ki))) for ki in range(5))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:38:20.650125Z",
     "iopub.status.busy": "2021-07-22T14:38:20.649649Z",
     "iopub.status.idle": "2021-07-22T14:38:20.660139Z",
     "shell.execute_reply": "2021-07-22T14:38:20.658951Z",
     "shell.execute_reply.started": "2021-07-22T14:38:20.650082Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: {'train': array([ True,  True,  True, ...,  True, False, False]),\n",
       "  'valid': array([False, False, False, ..., False,  True,  True])},\n",
       " 1: {'train': array([False,  True,  True, ...,  True,  True,  True]),\n",
       "  'valid': array([ True, False, False, ..., False, False, False])},\n",
       " 2: {'train': array([ True, False, False, ...,  True,  True,  True]),\n",
       "  'valid': array([False,  True,  True, ..., False, False, False])},\n",
       " 3: {'train': array([ True,  True,  True, ...,  True,  True,  True]),\n",
       "  'valid': array([False, False, False, ..., False, False, False])},\n",
       " 4: {'train': array([ True,  True,  True, ..., False,  True,  True]),\n",
       "  'valid': array([False, False, False, ...,  True, False, False])}}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_5 = split_k(train_df, 5)\n",
    "split_5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:43:55.738413Z",
     "iopub.status.busy": "2021-07-22T14:43:55.738038Z",
     "iopub.status.idle": "2021-07-22T14:43:55.746573Z",
     "shell.execute_reply": "2021-07-22T14:43:55.745359Z",
     "shell.execute_reply.started": "2021-07-22T14:43:55.738381Z"
    }
   },
   "outputs": [],
   "source": [
    "class LitDataMod(pl.LightningDataModule):\n",
    "    def __init__(self, train_df, test_df, split_map: Dict[str, np.ndarray]):\n",
    "        super().__init__()\n",
    "        self.train_df = train_df\n",
    "        self.test_df = test_df\n",
    "        self.split_map = split_map\n",
    "        \n",
    "    def train_dataloader(self):\n",
    "        df = self.train_df[self.split_map[\"train\"]]\n",
    "        return create_dl(df, istrain=True)\n",
    "    \n",
    "    def val_dataloader(self):\n",
    "        df = self.train_df[self.split_map[\"valid\"]]\n",
    "        return create_dl(df, istrain=False)\n",
    "    \n",
    "    def test_dataloader(self):\n",
    "        df = self.test_df\n",
    "        return create_dl(df, istrain=False)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:51:10.351444Z",
     "iopub.status.busy": "2021-07-22T14:51:10.350980Z",
     "iopub.status.idle": "2021-07-22T14:51:10.359417Z",
     "shell.execute_reply": "2021-07-22T14:51:10.358095Z",
     "shell.execute_reply.started": "2021-07-22T14:51:10.351410Z"
    }
   },
   "outputs": [],
   "source": [
    "model.pooler = nn.Linear(model.config.hidden_size, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Freeze/ unfreeze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ====================== GRADUAL UNFREEZE ======================\n",
    "def unfreeze(self):\n",
    "    \"\"\"unfreeze this module, and its sub modules\"\"\"\n",
    "    for p in self.parameters():\n",
    "        p.requires_grad = True\n",
    "        \n",
    "def freeze(self):\n",
    "    \"\"\"freeze this module, and its sub modules\"\"\"\n",
    "    for p in self.parameters():\n",
    "        p.requires_grad = False\n",
    "        \n",
    "def measure_freeze(m:nn.Module)->\"a describtion about how many submodules are unfreezed\":\n",
    "    total = 0\n",
    "    trainable = 0\n",
    "    for param in m.parameters():\n",
    "        total+=1\n",
    "        if param.requires_grad: trainable+=1\n",
    "    return f\"☀️ {trainable} trainables/{total} total\"\n",
    "        \n",
    "nn.Module.unfreeze = unfreeze\n",
    "nn.Module.freeze = freeze\n",
    "\n",
    "class GradualUnfreeze(pl.callbacks.Callback):\n",
    "    \"\"\"\n",
    "    A Callback that will deploy ULMFit's gradual unfreeze trick\n",
    "    \"\"\"\n",
    "    def on_epoch_start(self, trainer, pl_module):\n",
    "        epoch = trainer.current_epoch\n",
    "        \n",
    "        if epoch==0:\n",
    "            pl_module.freeze()\n",
    "        \n",
    "        # layers from top to bottom\n",
    "        glist = pl_module.gradual_list()\n",
    "        \n",
    "        if epoch<len(glist):\n",
    "            print(f\"🔥 unfreezing [{type(glist[epoch])}] on [epoch {epoch}]\")\n",
    "            glist[epoch].unfreeze()\n",
    "                \n",
    "            # print out statistics for current unfreezing\n",
    "            print(f\"{measure_freeze(pl_module)}\")\n",
    "            \n",
    "class FineTuneScheduler(pl.callbacks.Callback):\n",
    "    \"\"\"\n",
    "    Train the top layer for 2 epoch\n",
    "    then un freeze all\n",
    "    \"\"\"\n",
    "    def on_epoch_start(self, trainer, pl_module):\n",
    "        epoch = trainer.current_epoch\n",
    "        \n",
    "        if epoch==0:\n",
    "            pl_module.freeze()\n",
    "            for tl in pl_module.top_layers:\n",
    "                tl.unfreeze()\n",
    "        if epoch==CONFIG.frozen_years:\n",
    "            pl_module.unfreeze()\n",
    "            pl_module.base.embeddings.freeze()\n",
    "\n",
    "# ====================== OPTIMIZER CONFIGURATION ======================\n",
    "def modules_to_opt_conf(modules,**opt_kwargs)->\"a list of PyTorch optimizer config\":\n",
    "    \"\"\"\n",
    "    put in a sequence of pytorch modules, \n",
    "    return optimizer param groups of configs\n",
    "    \"\"\"\n",
    "    param_list = []\n",
    "    param_list_nd = []\n",
    "    no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
    "    for m in modules:\n",
    "        for n,p in m.named_parameters():\n",
    "            if any(nd in n for nd in no_decay):\n",
    "                param_list_nd.append(p)\n",
    "            else:\n",
    "                param_list.append(p)\n",
    "                \n",
    "    opt_kwargs_nd = deepcopy(opt_kwargs)\n",
    "    opt_kwargs_nd[\"weight_decay\"]=0.\n",
    "    return [dict(params=param_list,**opt_kwargs), # param_group with weight decay\n",
    "            dict(params=param_list_nd,**opt_kwargs_nd), # param_group without weight decay\n",
    "           ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_config(config, **kwargs):\n",
    "    for k, v in kwargs.items():\n",
    "        setattr(model.config, k, v)\n",
    "    return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "# inf_dl = LitDataMod(train_df, test_df, split_5[0]).test_dataloader()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:57:37.042125Z",
     "iopub.status.busy": "2021-07-22T14:57:37.041746Z",
     "iopub.status.idle": "2021-07-22T14:57:37.054399Z",
     "shell.execute_reply": "2021-07-22T14:57:37.053098Z",
     "shell.execute_reply.started": "2021-07-22T14:57:37.042094Z"
    }
   },
   "outputs": [],
   "source": [
    "class BaseModel(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.learning_rate = CONFIG.learning_rate\n",
    "        self.crit = nn.MSELoss()\n",
    "#         self.crit = nn.KLDivLoss()\n",
    "\n",
    "    def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int):\n",
    "        x = batch[\"excerpt\"]\n",
    "        y = batch[\"target\"]\n",
    "\n",
    "        y_ = self(x)\n",
    "        loss = self.crit(y_[:, 0], y)\n",
    "        self.log(\"loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int):\n",
    "        x = batch[\"excerpt\"]\n",
    "        y = batch[\"target\"]\n",
    "\n",
    "        y_ = self(x)\n",
    "        loss = self.crit(y_[:, 0], y)\n",
    "        self.log(\"val_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int):\n",
    "        x = batch[\"excerpt\"]\n",
    "        ids = batch[\"id\"]\n",
    "\n",
    "        y_ = self(x)\n",
    "\n",
    "        target = y_.detach().cpu().numpy().reshape(-1)\n",
    "\n",
    "        return pd.DataFrame({\"id\": ids, \"target\": target})\n",
    "\n",
    "    @property\n",
    "    def top_layers(self):\n",
    "        return [self.reg, ]\n",
    "\n",
    "    def gradual_list(self) -> List[nn.Module]:\n",
    "        \"\"\"\n",
    "        Return a list of list of nn.Modules, \n",
    "        for the modules to be unfreezed by the start of each epoch\n",
    "        \"\"\"\n",
    "        # top layer\n",
    "        rt = self.top_layers\n",
    "\n",
    "        # encoder layers\n",
    "        rt += list(self.base.encoder.layer)[::-1]\n",
    "\n",
    "        # embedding\n",
    "        rt += [model.embeddings, ]\n",
    "        return rt\n",
    "\n",
    "    def configure_optimizers(self,):\n",
    "        \"\"\"\n",
    "        Create the optimizer\n",
    "        \"\"\"\n",
    "        gradual = self.gradual_list()\n",
    "\n",
    "        no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
    "        h = self.config\n",
    "        # top layer param group\n",
    "        opt_configs = modules_to_opt_conf(self.top_layers,\n",
    "                                          lr=self.learning_rate*CONFIG.top_lr_expansion,\n",
    "                                          eps=h.adam_epsilon,\n",
    "                                          weight_decay=h.weight_decay)\n",
    "\n",
    "        delay = len(self.top_layers)\n",
    "        # encoder layer one by one from top to bottom param group\n",
    "        for i in range(len(gradual)-delay):\n",
    "            opt_configs += modules_to_opt_conf([gradual[i+delay], ],\n",
    "                                               lr=self.learning_rate * (h.layerwise_lr_decay**(i+1)),\n",
    "                                               eps=h.adam_epsilon,\n",
    "                                               weight_decay=h.weight_decay)\n",
    "\n",
    "        optimizer = torch.optim.AdamW(opt_configs)\n",
    "        self.opt = optimizer\n",
    "        return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Headers & Poolers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T14:57:37.042125Z",
     "iopub.status.busy": "2021-07-22T14:57:37.041746Z",
     "iopub.status.idle": "2021-07-22T14:57:37.054399Z",
     "shell.execute_reply": "2021-07-22T14:57:37.053098Z",
     "shell.execute_reply.started": "2021-07-22T14:57:37.042094Z"
    }
   },
   "outputs": [],
   "source": [
    "class CLSReg(BaseModel):\n",
    "    def __init__(self, base):\n",
    "        super().__init__()\n",
    "        self.base = base\n",
    "        self.config = self.base.config\n",
    "        self.reg = nn.Linear(base.config.hidden_size, 1)\n",
    "\n",
    "    def forward(self, excerpt: torch.Tensor) -> torch.Tensor:\n",
    "        x = excerpt[\"input_ids\"]\n",
    "        vec = self.base(x).last_hidden_state[:, 0, :]\n",
    "        return self.reg(vec)\n",
    "\n",
    "\n",
    "class CLSnn(BaseModel):\n",
    "    def __init__(self, base):\n",
    "        super().__init__()\n",
    "        self.base = base\n",
    "        self.config = self.base.config\n",
    "        self.reg = nn.Sequential(\n",
    "            nn.Linear(base.config.hidden_size,\n",
    "                      base.config.hidden_size, bias=False),\n",
    "            nn.BatchNorm1d(base.config.hidden_size),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(base.config.hidden_size, 1),\n",
    "        )\n",
    "\n",
    "    def forward(self, excerpt: torch.Tensor) -> torch.Tensor:\n",
    "        x = excerpt[\"input_ids\"]\n",
    "        vec = self.base(x).last_hidden_state[:, 0, :]\n",
    "        return self.reg(vec)\n",
    "    \n",
    "    \n",
    "class MeanPooler(BaseModel):\n",
    "    def __init__(self, base):\n",
    "        super().__init__()\n",
    "        self.base = base\n",
    "        self.config = self.base.config\n",
    "        self.reg = nn.Linear(base.config.hidden_size, 1)\n",
    "        self.dropout = nn.Dropout(CONFIG.drop_out_ratio)\n",
    "        \n",
    "    def reduce_hidden_state(self, vec, attention_mask):\n",
    "        vec = vec*(attention_mask[...,None])\n",
    "        return vec.sum(dim=1)/(attention_mask.sum(-1)[...,None])\n",
    "\n",
    "    def forward(self, excerpt: torch.Tensor) -> torch.Tensor:\n",
    "        x = excerpt[\"input_ids\"]\n",
    "        attention_mask = excerpt[\"attention_mask\"]\n",
    "        vec = self.base(x).last_hidden_state\n",
    "        vec = self.reduce_hidden_state(vec, attention_mask)\n",
    "        return self.reg(self.dropout(vec))\n",
    "\n",
    "\n",
    "# class MeanPooler(BaseModel):\n",
    "#     def __init__(self, base):\n",
    "#         super().__init__()\n",
    "#         self.base = base\n",
    "#         self.config = self.base.config\n",
    "#         self.reg = nn.Linear(base.config.hidden_size, 1)\n",
    "\n",
    "#     def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "#         vec = self.base(x).last_hidden_state.mean(dim=1)\n",
    "#         return self.reg(vec)\n",
    "\n",
    "\n",
    "class AttentionHead(nn.Module):\n",
    "    def __init__(self, in_features, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.in_features = in_features\n",
    "        self.middle_features = hidden_dim\n",
    "\n",
    "        self.W = nn.Linear(in_features, hidden_dim)\n",
    "        self.V = nn.Linear(hidden_dim, 1)\n",
    "        self.out_features = hidden_dim\n",
    "\n",
    "    def forward(self, features):\n",
    "        att = torch.tanh(self.W(features))\n",
    "        score = self.V(att)\n",
    "        attention_weights = torch.softmax(score, dim=1)\n",
    "        context_vector = attention_weights * features\n",
    "        context_vector = torch.sum(context_vector, dim=1)\n",
    "\n",
    "        return context_vector\n",
    "\n",
    "\n",
    "class WithAttnHead(BaseModel):\n",
    "    def __init__(self, base):\n",
    "        super().__init__()\n",
    "        self.base = base\n",
    "        self.config = self.base.config\n",
    "        self.head = AttentionHead(\n",
    "            self.config.hidden_size, self.config.hidden_size)\n",
    "        self.dout = nn.Dropout(.1)\n",
    "        self.reg = nn.Linear(base.config.hidden_size, 1)\n",
    "\n",
    "    def forward(self, excerpt: torch.Tensor) -> torch.Tensor:\n",
    "        x = excerpt[\"input_ids\"]\n",
    "        attention_mask = excerpt[\"attention_mask\"]\n",
    "        vec = self.base(x).last_hidden_state\n",
    "        vec *= (attention_mask[...,None])\n",
    "        return self.reg(self.dout(self.head(vec)))\n",
    "\n",
    "    @property\n",
    "    def top_layers(self):\n",
    "        return [self.reg, self.head]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Device and Save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> 2 cuda devices found >>>\n",
      "Device 0: \n",
      "\tname:Tesla V100-PCIE-32GB\n",
      "\tused:4MB\tfree:32506MB\n",
      "Device 1: \n",
      "\tname:Tesla V100-PCIE-32GB\n",
      "\tused:32003MB\tfree:507MB\n",
      "cuda stats refreshed\n",
      "Found the most idle GPU: cuda:0, 32506 MB Mem remained\n"
     ]
    }
   ],
   "source": [
    "dev = CudaHandler().idle().device.index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('weights/MeanPooler_ne_27_163525')"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ModelClass = MeanPooler\n",
    "\n",
    "train_time = datetime.now().strftime(\"%d_%H%M%S\")\n",
    "WEIGHTS = Path(f\"weights/{ModelClass.__name__}_ne_{train_time}/\")\n",
    "WEIGHTS.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "WEIGHTS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run folds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-22T15:07:07.774250Z",
     "iopub.status.busy": "2021-07-22T15:07:07.773521Z",
     "iopub.status.idle": "2021-07-22T15:07:07.782826Z",
     "shell.execute_reply": "2021-07-22T15:07:07.781437Z",
     "shell.execute_reply.started": "2021-07-22T15:07:07.774195Z"
    }
   },
   "outputs": [],
   "source": [
    "for fold in tqdm(range(len(split_5)), leave=False):\n",
    "    model = AutoModel.from_pretrained(pretrained)\n",
    "    lit_model = ModelClass(model)\n",
    "    lit_model.base.config = init_config(lit_model.base.config,\n",
    "        layerwise_lr_decay = 1.,\n",
    "        adam_epsilon = 1e-6,\n",
    "        weight_decay = 0.01\n",
    "    )\n",
    "#     lit_model.freeze()\n",
    "\n",
    "    saver = pl.callbacks.ModelCheckpoint(dirpath=WEIGHTS, monitor=\"val_loss\",mode=\"min\",\n",
    "                                         save_top_k=1, save_weights_only=True,\n",
    "                             filename='{epoch}-{val_loss:.2f}'+f\"fd{fold}\")\n",
    "    \n",
    "    early = pl.callbacks.EarlyStopping(\n",
    "        \"val_loss\", mode=\"min\", patience=5)\n",
    "    \n",
    "    trainer = pl.Trainer(\n",
    "        max_epochs=CONFIG.max_epochs,\n",
    "        callbacks=[saver, early, GradualUnfreeze()],\n",
    "        logger=False,\n",
    "        fast_dev_run=False,\n",
    "        gpus=[dev,],\n",
    "#         auto_lr_find = True,\n",
    "    )\n",
    "    trainer.tune(lit_model,datamodule=LitDataMod(\n",
    "        train_df, test_df, split_5[fold]))\n",
    "    print(f\"[LEARNING RATE]{lit_model.learning_rate}\")\n",
    "    trainer.fit(lit_model,datamodule=LitDataMod(\n",
    "        train_df, test_df, split_5[fold]))\n",
    "    lit_model = lit_model.cpu()\n",
    "\n",
    "    print(f\"[SAVE]{WEIGHTS}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deliver empety weights\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(WEIGHTS.ls()) ==0:\n",
    "    print(f\"[REMOVE EMPETY]{WEIGHTS}\")\n",
    "    WEIGHTS.rmdir()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}