notebooks/transformer_kfold.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# About the Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All the transformer models used for the final ensemble were based out of this notebook. This pipeline runs a 5-fold training. The folds are stratified using the iterstrat package which helps to stratify multi-label data. Specifically, `MultilabelStratifiedKFold` was used to create the folds. You can also check the complete data preprocessing and preparation stage in [this notebook](https://www.kaggle.com/code/atharvaingle/datasolve-eda-cv-setup)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "papermill": {
     "duration": 0.012442,
     "end_time": "2022-11-19T20:36:20.42823",
     "exception": false,
     "start_time": "2022-11-19T20:36:20.415788",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "papermill": {
     "duration": 0.031254,
     "end_time": "2022-11-19T20:36:20.46981",
     "exception": false,
     "start_time": "2022-11-19T20:36:20.438556",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "DEBUG = False\n",
    "WANDB = True\n",
    "ENVIRON = \"lambdalabs\"\n",
    "EXP_NAME = \"dbv3l-15ep\"\n",
    "NOTES = \"5 fold experiment, with the best settings till now, deberta-v3-large, 15 epochs, 4e-5, 16 bs, previous run failed because of catastrophic forgetting\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.009283,
     "end_time": "2022-11-19T20:36:20.488055",
     "exception": false,
     "start_time": "2022-11-19T20:36:20.478772",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Setup Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_kg_hide-output": true,
    "papermill": {
     "duration": 60.824307,
     "end_time": "2022-11-19T20:37:21.31941",
     "exception": false,
     "start_time": "2022-11-19T20:36:20.495103",
     "status": "completed"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pkgutil\n",
    "from pathlib import Path\n",
    "\n",
    "PROJECT = \"DataSolve-2022\"\n",
    "\n",
    "if ENVIRON == \"jarvislabs\":\n",
    "    ROOT_DIR = Path(f\"/home/{PROJECT}\")\n",
    "    ARTIFACTS_DIR = Path(\"/home/artifacts\")\n",
    "    SETUP_SCRIPT_PATH = Path(\"/home/setup.sh\")\n",
    "elif ENVIRON == \"lambdalabs\":\n",
    "    ROOT_DIR = Path(f\"/home/ubuntu/{PROJECT}\")\n",
    "    ARTIFACTS_DIR = Path(\"/home/ubuntu/artifacts\")\n",
    "    SETUP_SCRIPT_PATH = Path(\"/home/setup.sh\")\n",
    "elif ENVIRON == \"kaggle\":\n",
    "    ROOT_DIR = Path(f\"/kaggle/working/{PROJECT}\")\n",
    "    ARTIFACTS_DIR = Path(\"/kaggle/working/artifacts\")\n",
    "    SETUP_SCRIPT_PATH = Path(\"/kaggle/input/datasolve-setup-script/setup.sh\")\n",
    "    \n",
    "if not pkgutil.find_loader(\"omegaconf\") and ENVIRON == \"kaggle\":\n",
    "    !bash {SETUP_SCRIPT_PATH} {ENVIRON} \"true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "papermill": {
     "duration": 0.036363,
     "end_time": "2022-11-19T20:37:21.367118",
     "exception": false,
     "start_time": "2022-11-19T20:37:21.330755",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load secret keys\n",
    "%load_ext dotenv\n",
    "%dotenv {ROOT_DIR}/.env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.010899,
     "end_time": "2022-11-19T20:37:21.389126",
     "exception": false,
     "start_time": "2022-11-19T20:37:21.378227",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "papermill": {
     "duration": 0.092095,
     "end_time": "2022-11-19T20:37:21.492386",
     "exception": false,
     "start_time": "2022-11-19T20:37:21.400291",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving outputs to /home/ubuntu/artifacts/dbv3l-15ep\n",
      "EXPERIMENT: dbv3l-15ep, DESC: 5 fold experiment, with the best settings till now, deberta-v3-large, 15 epochs, 4e-5, 16 bs, previous run failed because of catastrophic forgetting\n",
      "\n",
      "debug: false\n",
      "wandb: true\n",
      "seed: 42\n",
      "train_csv: train_folds_5.csv\n",
      "fold: 0\n",
      "model:\n",
      "  model_name_or_path: microsoft/deberta-v3-large\n",
      "  gradient_checkpointing: false\n",
      "  reinit_last_layers: 0\n",
      "  output_hidden_states: false\n",
      "  output_last_hidden_state: false\n",
      "  output_pooled_embeds: false\n",
      "exp_name: dbv3l-15ep\n",
      "tags:\n",
      "- clspool\n",
      "- microsoft/deberta-v3-large\n",
      "- '512'\n",
      "- 5_fold_split\n",
      "notes: 5 fold experiment, with the best settings till now, deberta-v3-large, 15 epochs,\n",
      "  4e-5, 16 bs, previous run failed because of catastrophic forgetting\n",
      "upload_artifacts_to_wandb: true\n",
      "data:\n",
      "  max_length: 512\n",
      "  truncation: true\n",
      "  pad_to_multiple_of: 8\n",
      "training_args:\n",
      "  seed: 42\n",
      "  evaluation_strategy: epoch\n",
      "  save_strategy: epoch\n",
      "  save_total_limit: 1\n",
      "  num_train_epochs: 15\n",
      "  lr_scheduler_type: cosine\n",
      "  warmup_ratio: 0.2\n",
      "  per_device_train_batch_size: 16\n",
      "  per_device_eval_batch_size: 16\n",
      "  gradient_accumulation_steps: 1\n",
      "  learning_rate: 4.0e-05\n",
      "  weight_decay: 0.01\n",
      "  max_grad_norm: 1.0\n",
      "  adam_epsilon: 1.0e-06\n",
      "  fp16: true\n",
      "  dataloader_num_workers: 6\n",
      "  load_best_model_at_end: true\n",
      "  metric_for_best_model: eval_f1\n",
      "  greater_is_better: true\n",
      "  group_by_length: true\n",
      "  length_column_name: length\n",
      "  report_to: wandb\n",
      "  dataloader_pin_memory: true\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os, gc\n",
    "gc.enable()\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "class Config:\n",
    "    # GENERAL\n",
    "    debug = DEBUG\n",
    "    wandb = WANDB\n",
    "    seed = 42\n",
    "    train_csv = \"train_folds_5.csv\"\n",
    "    fold = 0 # will be overriden later\n",
    "    \n",
    "    # MODEL\n",
    "    model = dict(\n",
    "        model_name_or_path = \"microsoft/deberta-v3-large\",\n",
    "        gradient_checkpointing = False,\n",
    "        reinit_last_layers = 0,\n",
    "        output_hidden_states = False,\n",
    "        output_last_hidden_state = False,\n",
    "        output_pooled_embeds = False,\n",
    "    )\n",
    "\n",
    "\n",
    "    # TRACKING\n",
    "    exp_name = EXP_NAME\n",
    "    tags = [\"clspool\", f\"{model['model_name_or_path']}\", \"512\", \"5_fold_split\"]\n",
    "    notes = NOTES\n",
    "    upload_artifacts_to_wandb = True\n",
    "    \n",
    "    # DATA\n",
    "    data = dict(\n",
    "        max_length = 512,\n",
    "        truncation = True,\n",
    "        pad_to_multiple_of = 8,\n",
    "    )\n",
    "    \n",
    "    # TRAINING ARGUMENTS\n",
    "    training_args = dict(\n",
    "        # general\n",
    "        seed = seed,\n",
    "        evaluation_strategy = \"epoch\",\n",
    "        save_strategy = \"epoch\",\n",
    "        save_total_limit = 1,\n",
    "\n",
    "        # train settings\n",
    "        num_train_epochs = 15,\n",
    "        lr_scheduler_type = \"cosine\",\n",
    "        warmup_ratio = 0.2,\n",
    "        per_device_train_batch_size = 16,\n",
    "        per_device_eval_batch_size = 16,\n",
    "        gradient_accumulation_steps = 1,\n",
    "        learning_rate = 4e-5,\n",
    "        weight_decay = 0.01,\n",
    "        max_grad_norm = 1.0,\n",
    "        \n",
    "        # misc\n",
    "        adam_epsilon = 1e-6,\n",
    "        fp16 = True,\n",
    "        dataloader_num_workers = min(6, os.cpu_count()),\n",
    "        load_best_model_at_end = True,\n",
    "        metric_for_best_model = \"eval_f1\",\n",
    "        greater_is_better = True,\n",
    "        group_by_length = True,\n",
    "        length_column_name = \"length\",\n",
    "        report_to = \"wandb\" if WANDB else \"none\",\n",
    "        dataloader_pin_memory = True,\n",
    "    )\n",
    "\n",
    "\n",
    "# CONFIG SETTINGS\n",
    "config_dict = {x:dict(Config.__dict__)[x] for x in dict(Config.__dict__) if not x.startswith('_')}\n",
    "cfg = OmegaConf.create(config_dict)\n",
    "if cfg.debug: \n",
    "    cfg.tags += [\"debug\"]\n",
    "    cfg.training_args.num_train_epochs = 2\n",
    "    cfg.model.model_name_or_path = \"microsoft/deberta-v3-base\"\n",
    "\n",
    "OUTPUT_DIR = Path(ARTIFACTS_DIR/f'{cfg.exp_name}')\n",
    "print(f\"Saving outputs to {OUTPUT_DIR}\")\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "print(f\"EXPERIMENT: {cfg.exp_name}, DESC: {cfg.notes}\\n\")\n",
    "print(OmegaConf.to_yaml(cfg, resolve=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.010949,
     "end_time": "2022-11-19T20:37:21.514983",
     "exception": false,
     "start_time": "2022-11-19T20:37:21.504034",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "papermill": {
     "duration": 9.61634,
     "end_time": "2022-11-19T20:37:31.14242",
     "exception": false,
     "start_time": "2022-11-19T20:37:21.52608",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/.local/lib/python3.8/site-packages/pandas/core/computation/expressions.py:20: UserWarning: Pandas requires version '2.7.3' or newer of 'numexpr' (version '2.7.1' currently installed).\n",
      "  from pandas.core.computation.check import NUMEXPR_INSTALLED\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "import glob\n",
    "import shutil\n",
    "import pickle\n",
    "import warnings\n",
    "import logging\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pprint import pprint\n",
    "from dataclasses import dataclass\n",
    "from typing import Dict, List, Tuple, Callable, Optional, Union\n",
    "\n",
    "import wandb\n",
    "from wandb import AlertLevel\n",
    "\n",
    "from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit\n",
    "from sklearn.metrics import f1_score, accuracy_score, roc_auc_score\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import datasets, transformers\n",
    "from datasets import Dataset\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoConfig,\n",
    "    AutoModel,\n",
    "    AutoModelForSequenceClassification,\n",
    "    EvalPrediction,\n",
    "    PreTrainedTokenizer,\n",
    "    PretrainedConfig,\n",
    "    PreTrainedModel,\n",
    "    DataCollatorWithPadding,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.modeling_outputs import ModelOutput\n",
    "\n",
    "# SYSTEM SETTINGS\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "os.environ[\"TRANSFORMERS_NO_ADVISORY_WARNINGS\"] = \"true\"\n",
    "os.environ[\"WANDB_SILENT\"] = \"true\"\n",
    "set_seed(cfg.seed)\n",
    "if not cfg.debug:\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    logging.disable(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.011772,
     "end_time": "2022-11-19T20:37:31.166461",
     "exception": false,
     "start_time": "2022-11-19T20:37:31.154689",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "papermill": {
     "duration": 0.022708,
     "end_time": "2022-11-19T20:37:31.20094",
     "exception": false,
     "start_time": "2022-11-19T20:37:31.178232",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def delete_checkpoints(dir):\n",
    "    for file in glob.glob(f\"{dir}/checkpoint-*\"):\n",
    "        shutil.rmtree(file, ignore_errors=True)\n",
    "\n",
    "\n",
    "def clear_memory():\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "def delete_file(path: str):\n",
    "    if os.exists(path):\n",
    "        os.remove(path)\n",
    "\n",
    "def save_pickle(obj, filepath):\n",
    "    with open(filepath, 'wb') as handle:\n",
    "        pickle.dump(obj, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "    \n",
    "def process_config_for_wandb(cfg: OmegaConf):\n",
    "    \"\"\"\n",
    "    Only keep relevant part of config for logging\n",
    "    \"\"\"\n",
    "    tmp_cfg = copy.deepcopy(cfg)\n",
    "    cfg_dict = OmegaConf.to_container(tmp_cfg, resolve=True, throw_on_missing=True)\n",
    "    del cfg_dict[\"training_args\"]\n",
    "    return cfg_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.011751,
     "end_time": "2022-11-19T20:37:31.224629",
     "exception": false,
     "start_time": "2022-11-19T20:37:31.212878",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Read and process data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "papermill": {
     "duration": 0.252636,
     "end_time": "2022-11-19T20:37:31.489092",
     "exception": false,
     "start_time": "2022-11-19T20:37:31.236456",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50\n"
     ]
    },
    {
     "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>name</th>\n",
       "      <th>document_text</th>\n",
       "      <th>Accounting and Finance</th>\n",
       "      <th>Antitrust</th>\n",
       "      <th>Banking</th>\n",
       "      <th>Broker Dealer</th>\n",
       "      <th>Commodities Trading</th>\n",
       "      <th>Compliance Management</th>\n",
       "      <th>Consumer protection</th>\n",
       "      <th>...</th>\n",
       "      <th>Research</th>\n",
       "      <th>Risk Management</th>\n",
       "      <th>Securities Clearing</th>\n",
       "      <th>Securities Issuing</th>\n",
       "      <th>Securities Management</th>\n",
       "      <th>Securities Sales</th>\n",
       "      <th>Securities Settlement</th>\n",
       "      <th>Trade Pricing</th>\n",
       "      <th>Trade Settlement</th>\n",
       "      <th>fold</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4772</td>\n",
       "      <td>Consent Order in the Matter of Solium Financia...</td>\n",
       "      <td>Solium Financial Services LLC (\"SFS\") is a bro...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</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",
       "      <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>4774</td>\n",
       "      <td>Alberta Securities Commission Warns Investors ...</td>\n",
       "      <td>A new year brings new investment opportunities...</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",
       "      <td>0</td>\n",
       "      <td>...</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",
       "      <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>4775</td>\n",
       "      <td>Exempt Market Dealer Agrees to Settlement</td>\n",
       "      <td>The Alberta Securities Commission (ASC) has co...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4776</td>\n",
       "      <td>Canadian Securities Regulators Announces Consu...</td>\n",
       "      <td>The Canadian Securities Administrators (CSA) p...</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",
       "      <td>0</td>\n",
       "      <td>...</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",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4778</td>\n",
       "      <td>CSA Consultation Paper 51-405 Consideration of...</td>\n",
       "      <td>On April 6, 2017, the Canadian Securities Admi...</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",
       "      <td>1</td>\n",
       "      <td>...</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",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 54 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     id                                               name  \\\n",
       "0  4772  Consent Order in the Matter of Solium Financia...   \n",
       "1  4774  Alberta Securities Commission Warns Investors ...   \n",
       "2  4775          Exempt Market Dealer Agrees to Settlement   \n",
       "3  4776  Canadian Securities Regulators Announces Consu...   \n",
       "4  4778  CSA Consultation Paper 51-405 Consideration of...   \n",
       "\n",
       "                                       document_text  Accounting and Finance  \\\n",
       "0  Solium Financial Services LLC (\"SFS\") is a bro...                       0   \n",
       "1  A new year brings new investment opportunities...                       0   \n",
       "2  The Alberta Securities Commission (ASC) has co...                       0   \n",
       "3  The Canadian Securities Administrators (CSA) p...                       0   \n",
       "4  On April 6, 2017, the Canadian Securities Admi...                       0   \n",
       "\n",
       "   Antitrust  Banking  Broker Dealer  Commodities Trading  \\\n",
       "0          0        0              1                    0   \n",
       "1          0        0              0                    0   \n",
       "2          0        0              1                    0   \n",
       "3          0        0              0                    0   \n",
       "4          0        0              0                    0   \n",
       "\n",
       "   Compliance Management  Consumer protection  ...  Research  Risk Management  \\\n",
       "0                      1                    0  ...         0                0   \n",
       "1                      0                    0  ...         0                0   \n",
       "2                      1                    0  ...         0                0   \n",
       "3                      0                    0  ...         0                0   \n",
       "4                      0                    1  ...         0                0   \n",
       "\n",
       "   Securities Clearing  Securities Issuing  Securities Management  \\\n",
       "0                    0                   0                      0   \n",
       "1                    0                   0                      0   \n",
       "2                    0                   0                      0   \n",
       "3                    0                   0                      0   \n",
       "4                    0                   0                      0   \n",
       "\n",
       "   Securities Sales  Securities Settlement  Trade Pricing  Trade Settlement  \\\n",
       "0                 0                      0              0                 0   \n",
       "1                 0                      0              0                 0   \n",
       "2                 1                      1              0                 1   \n",
       "3                 0                      1              0                 0   \n",
       "4                 0                      1              0                 0   \n",
       "\n",
       "   fold  \n",
       "0     0  \n",
       "1     0  \n",
       "2     2  \n",
       "3     4  \n",
       "4     2  \n",
       "\n",
       "[5 rows x 54 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# READ DATA\n",
    "df = pd.read_csv(ROOT_DIR/'input'/cfg.train_csv)\n",
    "if cfg.debug:\n",
    "    df = df.sample(100, random_state=42).reset_index(drop=True)\n",
    "LABEL_COLS = [col for col in df.columns if col not in [\"id\", \"name\", \"document_text\", \"fold\"]]\n",
    "print(len(LABEL_COLS))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "papermill": {
     "duration": 24.305877,
     "end_time": "2022-11-19T20:37:56.092882",
     "exception": false,
     "start_time": "2022-11-19T20:37:31.787005",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   "
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b14495ef77004bd59885b1cd7851b246",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "#0:   0%|          | 0/4930 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f6bac773f4d48eda418b0d51b8b4da2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "#1:   0%|          | 0/4929 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(cfg.model.model_name_or_path)\n",
    "\n",
    "def tokenize_func(example: pd.Series, tokenizer: PreTrainedTokenizer, max_length: int = 512, truncation: bool = True, mode: str=\"train\"):\n",
    "    tokenized = tokenizer(\n",
    "    example[\"text\"],\n",
    "    truncation=truncation,\n",
    "    max_length=max_length,\n",
    "    add_special_tokens=True,\n",
    ")\n",
    "    if mode == \"train\":\n",
    "        tokenized[\"labels\"] = [example[i] for i in LABEL_COLS]\n",
    "    tokenized[\"length\"] = len(tokenized[\"input_ids\"])\n",
    "    return tokenized\n",
    "\n",
    "def preprocess_data(df_: pd.DataFrame, mode:str=\"train\"):\n",
    "    df_[\"text\"] = tokenizer.cls_token + df_[\"name\"] + tokenizer.sep_token + df_[\"document_text\"] + tokenizer.sep_token\n",
    "    tok_ds = Dataset.from_pandas(df_)\n",
    "    tok_ds = tok_ds.map(lambda x: tokenize_func(x, tokenizer, max_length=cfg.data.max_length, truncation=cfg.data.truncation, mode=mode), num_proc=2)\n",
    "    return tok_ds\n",
    "\n",
    "tok_ds = preprocess_data(df, mode=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "_kg_hide-output": true,
    "papermill": {
     "duration": 0.03757,
     "end_time": "2022-11-19T20:37:56.144422",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.106852",
     "status": "completed"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 4772,\n",
       " 'name': 'Consent Order in the Matter of Solium Financial Services LLC',\n",
       " 'document_text': 'Solium Financial Services LLC (\"SFS\") is a broker-dealer with a principal place of business at 50 Tice Boulevard, Suite A-18 Woodcliff Lake, New Jersey 07677, and is registered as a broker-dealer with the Alabama Securities Commission (\"Commission\"). During the period from at least January 2009 to June 6, 2019, SFS acted as broker-dealer in Alabama as the term broker-dealer is defined by Title 8, Chapter 6, 8-6-2 of the Act. Code of Alabama, 8-6-3(a) states that it is unlawful for a person to transact business in Alabama as a broker-dealer or agent unless such person is registered under the Act. By engaging in the conduct set forth above, SFS acted as an unregistered broker-dealer in Alabama in violation of 8-6-3(a) of the Act. This Order concludes the investigation by the Commission and any other action that the Commission could commence under applicable Alabama law as it relates to the substance of the Findings of Fact and Conclusions of Law herein, provided, however, that the Commission may pursue claims arising from SFS failure to comply with the terms of this Order. The fine, in this case, is $7,624.01.',\n",
       " 'Accounting and Finance': 0,\n",
       " 'Antitrust': 0,\n",
       " 'Banking': 0,\n",
       " 'Broker Dealer': 1,\n",
       " 'Commodities Trading': 0,\n",
       " 'Compliance Management': 1,\n",
       " 'Consumer protection': 0,\n",
       " 'Contract Provisions': 0,\n",
       " 'Corporate Communications': 0,\n",
       " 'Corporate Governance': 0,\n",
       " 'Definitions': 1,\n",
       " 'Delivery': 0,\n",
       " 'Examinations': 0,\n",
       " 'Exemptions': 0,\n",
       " 'Fees and Charges': 0,\n",
       " 'Financial Accounting': 0,\n",
       " 'Financial Crime': 0,\n",
       " 'Forms': 0,\n",
       " 'Fraud': 0,\n",
       " 'IT Risk': 0,\n",
       " 'Information Filing': 0,\n",
       " 'Insurance': 0,\n",
       " 'Legal': 1,\n",
       " 'Legal Proceedings': 1,\n",
       " 'Licensing': 1,\n",
       " 'Licensure and certification': 1,\n",
       " 'Liquidity Risk': 0,\n",
       " 'Listing': 0,\n",
       " 'Market Abuse': 0,\n",
       " 'Market Risk': 0,\n",
       " 'Monetary and Economic Policy': 0,\n",
       " 'Money Services': 0,\n",
       " 'Money-Laundering and Terrorist Financing': 0,\n",
       " 'Natural Disasters': 0,\n",
       " 'Payments and Settlements': 0,\n",
       " 'Powers and Duties': 0,\n",
       " 'Quotation': 0,\n",
       " 'Records Maintenance': 0,\n",
       " 'Regulatory Actions': 0,\n",
       " 'Regulatory Reporting': 0,\n",
       " 'Required Disclosures': 0,\n",
       " 'Research': 0,\n",
       " 'Risk Management': 0,\n",
       " 'Securities Clearing': 0,\n",
       " 'Securities Issuing': 0,\n",
       " 'Securities Management': 0,\n",
       " 'Securities Sales': 0,\n",
       " 'Securities Settlement': 0,\n",
       " 'Trade Pricing': 0,\n",
       " 'Trade Settlement': 0,\n",
       " 'fold': 0,\n",
       " 'text': '[CLS]Consent Order in the Matter of Solium Financial Services LLC[SEP]Solium Financial Services LLC (\"SFS\") is a broker-dealer with a principal place of business at 50 Tice Boulevard, Suite A-18 Woodcliff Lake, New Jersey 07677, and is registered as a broker-dealer with the Alabama Securities Commission (\"Commission\"). During the period from at least January 2009 to June 6, 2019, SFS acted as broker-dealer in Alabama as the term broker-dealer is defined by Title 8, Chapter 6, 8-6-2 of the Act. Code of Alabama, 8-6-3(a) states that it is unlawful for a person to transact business in Alabama as a broker-dealer or agent unless such person is registered under the Act. By engaging in the conduct set forth above, SFS acted as an unregistered broker-dealer in Alabama in violation of 8-6-3(a) of the Act. This Order concludes the investigation by the Commission and any other action that the Commission could commence under applicable Alabama law as it relates to the substance of the Findings of Fact and Conclusions of Law herein, provided, however, that the Commission may pursue claims arising from SFS failure to comply with the terms of this Order. The fine, in this case, is $7,624.01.[SEP]',\n",
       " 'input_ids': [1,\n",
       "  1,\n",
       "  36219,\n",
       "  4077,\n",
       "  267,\n",
       "  262,\n",
       "  14759,\n",
       "  265,\n",
       "  471,\n",
       "  60661,\n",
       "  3729,\n",
       "  1724,\n",
       "  3927,\n",
       "  2,\n",
       "  471,\n",
       "  60661,\n",
       "  3729,\n",
       "  1724,\n",
       "  3927,\n",
       "  287,\n",
       "  309,\n",
       "  430,\n",
       "  16480,\n",
       "  309,\n",
       "  285,\n",
       "  269,\n",
       "  266,\n",
       "  7347,\n",
       "  271,\n",
       "  58161,\n",
       "  275,\n",
       "  266,\n",
       "  4891,\n",
       "  470,\n",
       "  265,\n",
       "  460,\n",
       "  288,\n",
       "  960,\n",
       "  897,\n",
       "  10953,\n",
       "  14201,\n",
       "  261,\n",
       "  9500,\n",
       "  336,\n",
       "  271,\n",
       "  2048,\n",
       "  4059,\n",
       "  40099,\n",
       "  2202,\n",
       "  261,\n",
       "  485,\n",
       "  3744,\n",
       "  7844,\n",
       "  46692,\n",
       "  261,\n",
       "  263,\n",
       "  269,\n",
       "  2079,\n",
       "  283,\n",
       "  266,\n",
       "  7347,\n",
       "  271,\n",
       "  58161,\n",
       "  275,\n",
       "  262,\n",
       "  6002,\n",
       "  10207,\n",
       "  2653,\n",
       "  287,\n",
       "  309,\n",
       "  75573,\n",
       "  309,\n",
       "  285,\n",
       "  260,\n",
       "  1717,\n",
       "  262,\n",
       "  926,\n",
       "  292,\n",
       "  288,\n",
       "  668,\n",
       "  1278,\n",
       "  1812,\n",
       "  264,\n",
       "  1172,\n",
       "  525,\n",
       "  261,\n",
       "  1112,\n",
       "  261,\n",
       "  10000,\n",
       "  430,\n",
       "  8736,\n",
       "  283,\n",
       "  7347,\n",
       "  271,\n",
       "  58161,\n",
       "  267,\n",
       "  6002,\n",
       "  283,\n",
       "  262,\n",
       "  1384,\n",
       "  7347,\n",
       "  271,\n",
       "  58161,\n",
       "  269,\n",
       "  3034,\n",
       "  293,\n",
       "  7181,\n",
       "  578,\n",
       "  261,\n",
       "  4696,\n",
       "  525,\n",
       "  261,\n",
       "  578,\n",
       "  271,\n",
       "  765,\n",
       "  271,\n",
       "  445,\n",
       "  265,\n",
       "  262,\n",
       "  1878,\n",
       "  260,\n",
       "  3506,\n",
       "  265,\n",
       "  6002,\n",
       "  261,\n",
       "  578,\n",
       "  271,\n",
       "  765,\n",
       "  271,\n",
       "  508,\n",
       "  555,\n",
       "  452,\n",
       "  285,\n",
       "  1603,\n",
       "  272,\n",
       "  278,\n",
       "  269,\n",
       "  15082,\n",
       "  270,\n",
       "  266,\n",
       "  604,\n",
       "  264,\n",
       "  51494,\n",
       "  460,\n",
       "  267,\n",
       "  6002,\n",
       "  283,\n",
       "  266,\n",
       "  7347,\n",
       "  271,\n",
       "  58161,\n",
       "  289,\n",
       "  2645,\n",
       "  2336,\n",
       "  405,\n",
       "  604,\n",
       "  269,\n",
       "  2079,\n",
       "  494,\n",
       "  262,\n",
       "  1878,\n",
       "  260,\n",
       "  927,\n",
       "  4686,\n",
       "  267,\n",
       "  262,\n",
       "  3360,\n",
       "  487,\n",
       "  4243,\n",
       "  764,\n",
       "  261,\n",
       "  10000,\n",
       "  430,\n",
       "  8736,\n",
       "  283,\n",
       "  299,\n",
       "  46245,\n",
       "  7347,\n",
       "  271,\n",
       "  58161,\n",
       "  267,\n",
       "  6002,\n",
       "  267,\n",
       "  6527,\n",
       "  265,\n",
       "  578,\n",
       "  271,\n",
       "  765,\n",
       "  271,\n",
       "  508,\n",
       "  555,\n",
       "  452,\n",
       "  285,\n",
       "  265,\n",
       "  262,\n",
       "  1878,\n",
       "  260,\n",
       "  329,\n",
       "  4077,\n",
       "  14207,\n",
       "  262,\n",
       "  2485,\n",
       "  293,\n",
       "  262,\n",
       "  2653,\n",
       "  263,\n",
       "  356,\n",
       "  340,\n",
       "  1016,\n",
       "  272,\n",
       "  262,\n",
       "  2653,\n",
       "  387,\n",
       "  15851,\n",
       "  494,\n",
       "  4190,\n",
       "  6002,\n",
       "  818,\n",
       "  283,\n",
       "  278,\n",
       "  9611,\n",
       "  264,\n",
       "  262,\n",
       "  5182,\n",
       "  265,\n",
       "  262,\n",
       "  38269,\n",
       "  265,\n",
       "  17696,\n",
       "  263,\n",
       "  42068,\n",
       "  265,\n",
       "  2321,\n",
       "  10349,\n",
       "  261,\n",
       "  949,\n",
       "  261,\n",
       "  901,\n",
       "  261,\n",
       "  272,\n",
       "  262,\n",
       "  2653,\n",
       "  372,\n",
       "  5012,\n",
       "  2071,\n",
       "  9264,\n",
       "  292,\n",
       "  10000,\n",
       "  430,\n",
       "  2694,\n",
       "  264,\n",
       "  6378,\n",
       "  275,\n",
       "  262,\n",
       "  1169,\n",
       "  265,\n",
       "  291,\n",
       "  4077,\n",
       "  260,\n",
       "  279,\n",
       "  1399,\n",
       "  261,\n",
       "  267,\n",
       "  291,\n",
       "  571,\n",
       "  261,\n",
       "  269,\n",
       "  419,\n",
       "  819,\n",
       "  261,\n",
       "  48328,\n",
       "  260,\n",
       "  3085,\n",
       "  260,\n",
       "  2,\n",
       "  2],\n",
       " 'token_type_ids': [0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0],\n",
       " 'attention_mask': [1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1],\n",
       " 'labels': [0,\n",
       "  0,\n",
       "  0,\n",
       "  1,\n",
       "  0,\n",
       "  1,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  1,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0],\n",
       " 'length': 276}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tok_ds[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.013486,
     "end_time": "2022-11-19T20:37:56.1719",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.158414",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "papermill": {
     "duration": 0.024145,
     "end_time": "2022-11-19T20:37:56.209778",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.185633",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def post_process_logits(logits: np.ndarray, threshold=0.5):\n",
    "    # first, apply sigmoid on logits which are of shape (batch_size, num_labels)\n",
    "    sigmoid = torch.nn.Sigmoid()\n",
    "    probs = sigmoid(torch.Tensor(logits))\n",
    "    # next, use threshold to turn them into integer predictions\n",
    "    preds = np.zeros(probs.shape)\n",
    "    preds[np.where(probs >= threshold)] = 1\n",
    "    preds = preds.flatten().astype(int)\n",
    "    return preds\n",
    "\n",
    "def compute_metrics(p: EvalPrediction):\n",
    "    # `predictions` might return last_hidden_state or pooled_embeds\n",
    "    # In that case, take the first element (array) of the tuple for logits\n",
    "    logits = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "    preds = post_process_logits(logits)\n",
    "    labels = p.label_ids.flatten()\n",
    "    f1_macro_average = f1_score(labels, preds, average='macro')\n",
    "    roc_auc = roc_auc_score(labels, preds, average = 'macro')\n",
    "    accuracy = accuracy_score(labels, preds)\n",
    "    # return as dictionary\n",
    "    return {\n",
    "        'f1': f1_macro_average,\n",
    "        'roc_auc': roc_auc,\n",
    "        'accuracy': accuracy\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.013594,
     "end_time": "2022-11-19T20:37:56.237049",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.223455",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Custom Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "papermill": {
     "duration": 0.021072,
     "end_time": "2022-11-19T20:37:56.271589",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.250517",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# class MultilabelTrainer(Trainer):\n",
    "#     def compute_loss(self, model, inputs, return_outputs=False):\n",
    "#         labels = inputs.pop(\"labels\")\n",
    "#         outputs = model(**inputs)\n",
    "#         logits = outputs.logits\n",
    "#         loss_fct = torch.nn.BCEWithLogitsLoss()\n",
    "#         loss = loss_fct(logits.view(-1, self.model.config.num_labels), \n",
    "#                         labels.float().view(-1, self.model.config.num_labels))\n",
    "#         return (loss, outputs) if return_outputs else loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.013631,
     "end_time": "2022-11-19T20:37:56.298921",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.28529",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "papermill": {
     "duration": 0.0286,
     "end_time": "2022-11-19T20:37:56.341355",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.312755",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class CustomModelOutput(ModelOutput):\n",
    "    loss: Optional[torch.FloatTensor] = None\n",
    "    logits: torch.FloatTensor = None\n",
    "    hidden_states: Optional[Tuple[torch.FloatTensor]] = None\n",
    "    last_hidden_state: Optional[torch.FloatTensor] = None\n",
    "    pooled_embeds: Optional[torch.FloatTensor] = None\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ModelConfig:\n",
    "    model_name_or_path: str\n",
    "    gradient_checkpointing: Optional[bool] = False\n",
    "    reinit_last_layers: Optional[int] = 0\n",
    "    output_hidden_states: Optional[bool] = False\n",
    "    output_last_hidden_state: Optional[bool] = False\n",
    "    output_pooled_embeds: Optional[bool] = False\n",
    "\n",
    "def reinit_last_layers(model: Union[nn.Module, PreTrainedModel], num_layers: int):\n",
    "    \"\"\"Re-initialize the last-k transformer layers\"\"\"\n",
    "    if num_layers > 0:\n",
    "        model.encoder.layer[-num_layers:].apply(model._init_weights)\n",
    "    \n",
    "class DataSolveModel(nn.Module):\n",
    "    def __init__(self, config: ModelConfig):\n",
    "        super().__init__()\n",
    "        self.model_config = config\n",
    "        self.hf_config = AutoConfig.from_pretrained(self.model_config.model_name_or_path)\n",
    "        \n",
    "        self.backbone = AutoModel.from_pretrained(self.model_config.model_name_or_path, config=self.hf_config)\n",
    "        \n",
    "        if self.model_config.gradient_checkpointing:\n",
    "            self.backbone.gradient_checkpointing_enable()\n",
    "\n",
    "        # Initialize last-k transformer (backbone) layers\n",
    "        reinit_last_layers(self.backbone, self.model_config.reinit_last_layers)\n",
    "        \n",
    "        self.output = nn.Linear(self.hf_config.hidden_size, 50)\n",
    "        \n",
    "    def forward(self, input_ids, attention_mask, labels=None):\n",
    "        trans_out = self.backbone(input_ids, attention_mask=attention_mask)\n",
    "        last_hidden_state = trans_out.last_hidden_state\n",
    "        pooled_embeds = last_hidden_state[:, 0] # CLS Token\n",
    "        \n",
    "        logits = self.output(pooled_embeds)\n",
    "\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss_fct = torch.nn.BCEWithLogitsLoss()\n",
    "            loss = loss_fct(logits.view(-1, 50), labels.float().view(-1, 50))\n",
    "            \n",
    "        return CustomModelOutput(\n",
    "            loss=loss,\n",
    "            logits=logits,\n",
    "            hidden_states=trans_out.hidden_states if self.model_config.output_hidden_states else None,\n",
    "            last_hidden_state=last_hidden_state if self.model_config.output_last_hidden_state else None,\n",
    "            pooled_embeds=pooled_embeds if self.model_config.output_pooled_embeds else None,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.013329,
     "end_time": "2022-11-19T20:37:56.368443",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.355114",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "papermill": {
     "duration": 15677.599997,
     "end_time": "2022-11-20T00:59:13.982221",
     "exception": false,
     "start_time": "2022-11-19T20:37:56.382224",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_fold(cfg, fold):\n",
    "    set_seed(cfg.seed)\n",
    "    cfg.fold = fold\n",
    "    if cfg.wandb:\n",
    "        # init wandb run\n",
    "        wandb.init(\n",
    "            project=\"DataSolve-2022\",\n",
    "            group=cfg.exp_name,\n",
    "            name=f\"fold_{cfg.fold}\",\n",
    "            tags=cfg.tags,\n",
    "            notes=cfg.notes,\n",
    "            config=process_config_for_wandb(cfg),\n",
    "            save_code=True,\n",
    "        )\n",
    "        # send slack notification\n",
    "        wandb.alert(\n",
    "            title=f\"Experiment {cfg.exp_name}\",\n",
    "            text=f\"🚀 Starting experiment {cfg.exp_name} (fold_{cfg.fold}), Description: {cfg.notes}\",\n",
    "            level=AlertLevel.INFO,\n",
    "            wait_duration=0,\n",
    "        )\n",
    "\n",
    "    OUT_DIR = OUTPUT_DIR/f'fold_{cfg.fold}'\n",
    "    os.makedirs(OUT_DIR, exist_ok=True)\n",
    "    \n",
    "    # filter train and val data\n",
    "    train_ds = tok_ds.filter(lambda x: x[\"fold\"] != cfg.fold, desc=\"Filtering train idxs\")\n",
    "    val_ds = tok_ds.filter(lambda x: x[\"fold\"] == cfg.fold, desc=\"Filtering valid idxs\")\n",
    "\n",
    "    # sort by length to have similar length samples in each batch for speeding up evaluation\n",
    "    val_ds = val_ds.sort(\"length\")\n",
    "    val_ids = val_ds['id']\n",
    "    # remove unwanted columns\n",
    "    keep_cols = {\"input_ids\", \"attention_mask\", \"labels\", \"token_type_ids\"}\n",
    "    remove_cols = [c for c in train_ds.column_names if c not in keep_cols]\n",
    "    train_ds = train_ds.remove_columns(remove_cols)\n",
    "    val_ds = val_ds.remove_columns(remove_cols)\n",
    "    train_ds.set_format(\"torch\")\n",
    "    val_ds.set_format(\"torch\")\n",
    "\n",
    "    # init model\n",
    "    model_config = ModelConfig(**cfg.model)\n",
    "    model = DataSolveModel(model_config)\n",
    "\n",
    "    # init trainer\n",
    "    training_args = TrainingArguments(output_dir=OUT_DIR, **cfg.training_args)\n",
    "    trainer = Trainer(\n",
    "                model,\n",
    "                args=training_args,\n",
    "                data_collator=DataCollatorWithPadding(tokenizer, pad_to_multiple_of=cfg.data.pad_to_multiple_of),\n",
    "                train_dataset=train_ds,\n",
    "                eval_dataset=val_ds,\n",
    "                tokenizer=tokenizer,\n",
    "                compute_metrics=compute_metrics,\n",
    "    )\n",
    "    # train\n",
    "    trainer.train()\n",
    "\n",
    "    # ---------------------------------- Save, log, cleanup and upload to W&B ------------------------\n",
    "\n",
    "    # Save model\n",
    "    delete_checkpoints(OUT_DIR)\n",
    "    trainer.save_model()\n",
    "    clear_memory()\n",
    "\n",
    "    # Infer on validation set and extract logits and labels\n",
    "    eval_out = trainer.predict(val_ds)\n",
    "    logits = eval_out.predictions[0] if isinstance(eval_out.predictions, tuple) else eval_out.predictions\n",
    "    labels = eval_out.label_ids\n",
    "    \n",
    "    oof_dict = {\"id\": val_ids, \"logits\": logits, \"labels\": labels}\n",
    "    save_pickle(oof_dict, OUT_DIR/f\"oof_{cfg.fold}.pkl\")\n",
    "    \n",
    "    fin_f1_score = np.round(eval_out.metrics[\"test_f1\"], 6)\n",
    "\n",
    "    print(\"*\" * 30)\n",
    "    print(f\"  Experiment {cfg.exp_name}, Fold {cfg.fold}, F1-SCORE: {fin_f1_score}\")\n",
    "    print(\"*\" * 30)\n",
    "\n",
    "    # save experiment config file\n",
    "    config_file_save_path = OUT_DIR / f\"{cfg.exp_name}_config.yaml\"\n",
    "    with open(config_file_save_path, \"w\") as fp:\n",
    "        OmegaConf.save(config=cfg, f=fp.name)\n",
    "\n",
    "    # log artifacts to wandb\n",
    "    if cfg.wandb:\n",
    "        wandb.log({\"cv\": fin_f1_score})\n",
    "        if cfg.upload_artifacts_to_wandb:\n",
    "            shutil.copyfile(config_file_save_path, os.path.join(wandb.run.dir, f\"{cfg.exp_name}_config.yaml\"))\n",
    "            model_artifact = wandb.Artifact(name=cfg.exp_name, type=\"model\")\n",
    "            model_artifact.add_dir(OUT_DIR)\n",
    "            wandb.log_artifact(model_artifact, aliases=f\"fold_{cfg.fold}\")\n",
    "\n",
    "        wandb.alert(\n",
    "            title=f\"Experiment {cfg.exp_name}\",\n",
    "            text=f\"🎉 Finished experiment {cfg.exp_name} (fold_{cfg.fold}), Score: {fin_f1_score}\",\n",
    "            level=AlertLevel.INFO,\n",
    "            wait_duration=0,\n",
    "        )\n",
    "        wandb.finish()\n",
    "            \n",
    "    del model, trainer, eval_out, train_ds, val_ds; clear_memory();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "19fe74cb534648eaa48f5de189f6eee5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filtering train idxs:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ca30f24a1a14cd7bd77a6078dcddf75",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filtering valid idxs:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='7410' max='7410' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [7410/7410 44:12, Epoch 15/15]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>F1</th>\n",
       "      <th>Roc Auc</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.235258</td>\n",
       "      <td>0.648483</td>\n",
       "      <td>0.605215</td>\n",
       "      <td>0.919256</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.372900</td>\n",
       "      <td>0.190268</td>\n",
       "      <td>0.718239</td>\n",
       "      <td>0.657466</td>\n",
       "      <td>0.931961</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.206300</td>\n",
       "      <td>0.147810</td>\n",
       "      <td>0.820948</td>\n",
       "      <td>0.767850</td>\n",
       "      <td>0.948467</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.160500</td>\n",
       "      <td>0.124416</td>\n",
       "      <td>0.860252</td>\n",
       "      <td>0.818393</td>\n",
       "      <td>0.957229</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.126800</td>\n",
       "      <td>0.106454</td>\n",
       "      <td>0.886075</td>\n",
       "      <td>0.852261</td>\n",
       "      <td>0.963953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.097800</td>\n",
       "      <td>0.098110</td>\n",
       "      <td>0.897677</td>\n",
       "      <td>0.868193</td>\n",
       "      <td>0.967132</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.074200</td>\n",
       "      <td>0.093254</td>\n",
       "      <td>0.904980</td>\n",
       "      <td>0.876075</td>\n",
       "      <td>0.969384</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.055900</td>\n",
       "      <td>0.087396</td>\n",
       "      <td>0.917118</td>\n",
       "      <td>0.910615</td>\n",
       "      <td>0.971666</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.040800</td>\n",
       "      <td>0.081806</td>\n",
       "      <td>0.923885</td>\n",
       "      <td>0.916153</td>\n",
       "      <td>0.974050</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.029100</td>\n",
       "      <td>0.080604</td>\n",
       "      <td>0.928175</td>\n",
       "      <td>0.917849</td>\n",
       "      <td>0.975660</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.020700</td>\n",
       "      <td>0.080391</td>\n",
       "      <td>0.928271</td>\n",
       "      <td>0.920762</td>\n",
       "      <td>0.975527</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.015500</td>\n",
       "      <td>0.080391</td>\n",
       "      <td>0.929741</td>\n",
       "      <td>0.923143</td>\n",
       "      <td>0.975976</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>0.012600</td>\n",
       "      <td>0.080176</td>\n",
       "      <td>0.929649</td>\n",
       "      <td>0.921607</td>\n",
       "      <td>0.976026</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>0.010800</td>\n",
       "      <td>0.080252</td>\n",
       "      <td>0.930153</td>\n",
       "      <td>0.921898</td>\n",
       "      <td>0.976210</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>0.010100</td>\n",
       "      <td>0.080361</td>\n",
       "      <td>0.930020</td>\n",
       "      <td>0.921870</td>\n",
       "      <td>0.976159</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "******************************\n",
      "  Experiment dbv3l-15ep, Fold 0, F1-SCORE: 0.930153\n",
      "******************************\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6b8b978aa1e74e919efb9a5b656b39dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.01666862536667395, max=1.0)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29700f34e7c94902837b0dc903d2c609",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filtering train idxs:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43c44188284f4655b66c5dc1c6b5f1fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filtering valid idxs:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='7380' max='7380' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [7380/7380 44:13, Epoch 15/15]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>F1</th>\n",
       "      <th>Roc Auc</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.234580</td>\n",
       "      <td>0.595673</td>\n",
       "      <td>0.567181</td>\n",
       "      <td>0.916250</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.372100</td>\n",
       "      <td>0.176830</td>\n",
       "      <td>0.772997</td>\n",
       "      <td>0.718693</td>\n",
       "      <td>0.938740</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.205000</td>\n",
       "      <td>0.152576</td>\n",
       "      <td>0.802578</td>\n",
       "      <td>0.745916</td>\n",
       "      <td>0.945800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.159900</td>\n",
       "      <td>0.121493</td>\n",
       "      <td>0.867104</td>\n",
       "      <td>0.836034</td>\n",
       "      <td>0.958570</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.124800</td>\n",
       "      <td>0.106480</td>\n",
       "      <td>0.881656</td>\n",
       "      <td>0.844629</td>\n",
       "      <td>0.963620</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.096100</td>\n",
       "      <td>0.099585</td>\n",
       "      <td>0.891668</td>\n",
       "      <td>0.856098</td>\n",
       "      <td>0.966460</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.073400</td>\n",
       "      <td>0.088934</td>\n",
       "      <td>0.911911</td>\n",
       "      <td>0.894989</td>\n",
       "      <td>0.971200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.055200</td>\n",
       "      <td>0.084161</td>\n",
       "      <td>0.919702</td>\n",
       "      <td>0.909699</td>\n",
       "      <td>0.973270</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.039400</td>\n",
       "      <td>0.084855</td>\n",
       "      <td>0.920890</td>\n",
       "      <td>0.918086</td>\n",
       "      <td>0.973200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.027700</td>\n",
       "      <td>0.080498</td>\n",
       "      <td>0.928100</td>\n",
       "      <td>0.920766</td>\n",
       "      <td>0.975900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.019400</td>\n",
       "      <td>0.079403</td>\n",
       "      <td>0.929003</td>\n",
       "      <td>0.920817</td>\n",
       "      <td>0.976250</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.014600</td>\n",
       "      <td>0.080462</td>\n",
       "      <td>0.929537</td>\n",
       "      <td>0.925592</td>\n",
       "      <td>0.976190</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>0.011900</td>\n",
       "      <td>0.079827</td>\n",
       "      <td>0.930361</td>\n",
       "      <td>0.922947</td>\n",
       "      <td>0.976660</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>0.010400</td>\n",
       "      <td>0.080004</td>\n",
       "      <td>0.930625</td>\n",
       "      <td>0.923002</td>\n",
       "      <td>0.976760</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>0.009800</td>\n",
       "      <td>0.080025</td>\n",
       "      <td>0.930679</td>\n",
       "      <td>0.923198</td>\n",
       "      <td>0.976770</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "******************************\n",
      "  Experiment dbv3l-15ep, Fold 1, F1-SCORE: 0.930679\n",
      "******************************\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e109ae275df74ae9bdcb36223d98887b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.016669114000008752, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42d1f7bd18094c38b2e1121b69eb958a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filtering train idxs:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1565bf9a53574bb5acacee6391048348",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filtering valid idxs:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='6362' max='7410' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [6362/7410 37:59 < 06:15, 2.79 it/s, Epoch 12.88/15]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>F1</th>\n",
       "      <th>Roc Auc</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.232452</td>\n",
       "      <td>0.636522</td>\n",
       "      <td>0.595654</td>\n",
       "      <td>0.918852</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.371000</td>\n",
       "      <td>0.181365</td>\n",
       "      <td>0.762802</td>\n",
       "      <td>0.702579</td>\n",
       "      <td>0.937815</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.205500</td>\n",
       "      <td>0.147619</td>\n",
       "      <td>0.827435</td>\n",
       "      <td>0.781240</td>\n",
       "      <td>0.948902</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.157900</td>\n",
       "      <td>0.128414</td>\n",
       "      <td>0.843465</td>\n",
       "      <td>0.790447</td>\n",
       "      <td>0.954400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.127300</td>\n",
       "      <td>0.110345</td>\n",
       "      <td>0.873745</td>\n",
       "      <td>0.828483</td>\n",
       "      <td>0.961677</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.097900</td>\n",
       "      <td>0.100820</td>\n",
       "      <td>0.890632</td>\n",
       "      <td>0.851090</td>\n",
       "      <td>0.966016</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.075300</td>\n",
       "      <td>0.092039</td>\n",
       "      <td>0.909216</td>\n",
       "      <td>0.892213</td>\n",
       "      <td>0.969868</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.056700</td>\n",
       "      <td>0.086202</td>\n",
       "      <td>0.919435</td>\n",
       "      <td>0.909831</td>\n",
       "      <td>0.972734</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.040700</td>\n",
       "      <td>0.081471</td>\n",
       "      <td>0.926290</td>\n",
       "      <td>0.909682</td>\n",
       "      <td>0.975467</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.028700</td>\n",
       "      <td>0.081438</td>\n",
       "      <td>0.927258</td>\n",
       "      <td>0.911583</td>\n",
       "      <td>0.975732</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.020500</td>\n",
       "      <td>0.081155</td>\n",
       "      <td>0.928103</td>\n",
       "      <td>0.918238</td>\n",
       "      <td>0.975671</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.015500</td>\n",
       "      <td>0.080887</td>\n",
       "      <td>0.928733</td>\n",
       "      <td>0.918326</td>\n",
       "      <td>0.975915</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for fold in range(5):\n",
    "    train_fold(cfg, fold)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.018005,
     "end_time": "2022-11-20T00:59:14.021706",
     "exception": false,
     "start_time": "2022-11-20T00:59:14.003701",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init a new wandb run for storing submission artifacts\n",
    "wandb.init(\n",
    "    project=\"DataSolve-2022\",\n",
    "    group=cfg.exp_name,\n",
    "    name=\"inference\",\n",
    "    tags=cfg.tags,\n",
    "    notes=cfg.notes,\n",
    "    config=process_config_for_wandb(cfg)\n",
    ")\n",
    "SUB_OUT_DIR = OUTPUT_DIR/'submission'\n",
    "os.makedirs(SUB_OUT_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 1.215633,
     "end_time": "2022-11-20T00:59:15.258915",
     "exception": false,
     "start_time": "2022-11-20T00:59:14.043282",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_df = pd.read_csv(ROOT_DIR/'input'/'test.csv')\n",
    "sub_df = pd.read_csv(ROOT_DIR/'input'/'sample_submission.csv')\n",
    "test_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 2.0626,
     "end_time": "2022-11-20T00:59:17.336512",
     "exception": false,
     "start_time": "2022-11-20T00:59:15.273912",
     "status": "completed"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_ds = preprocess_data(test_df, mode=\"inference\")\n",
    "# sort test dataset to have similar length samples in a batch to speed up inference\n",
    "test_ds = test_ds.sort(\"length\")\n",
    "test_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 183.3963,
     "end_time": "2022-11-20T01:02:29.119814",
     "exception": false,
     "start_time": "2022-11-20T00:59:25.723514",
     "status": "completed"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_logits = []\n",
    "for fold in range(5):\n",
    "    print(f\"{'='*10} Inferring FOLD-{fold} {'='*10}\")\n",
    "    model_dir = OUTPUT_DIR/f'fold_{fold}'\n",
    "    model_config = ModelConfig(**cfg.model)\n",
    "    model = DataSolveModel(model_config)\n",
    "    model.load_state_dict(torch.load(model_dir/'pytorch_model.bin'))\n",
    "    trainer_args = TrainingArguments(\"./tmp\", per_device_eval_batch_size = 16)\n",
    "    trainer = Trainer(model, trainer_args, data_collator=DataCollatorWithPadding(tokenizer))\n",
    "    out = trainer.predict(test_ds)\n",
    "    logits = out.predictions[0] if isinstance(out.predictions, tuple) else out.predictions\n",
    "    all_logits.append(logits)\n",
    "    \n",
    "    del model, trainer; clear_memory();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fin_logits = np.mean(all_logits, axis=0)\n",
    "# save test logits to a file\n",
    "test_dict = {\"id\": test_ds['id'], \"logits\": fin_logits}\n",
    "save_pickle(test_dict, SUB_OUT_DIR/f\"{cfg.exp_name}_test_logits.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.014672,
     "end_time": "2022-11-20T01:02:29.14982",
     "exception": false,
     "start_time": "2022-11-20T01:02:29.135148",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Create submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 0.868639,
     "end_time": "2022-11-20T01:02:30.033557",
     "exception": false,
     "start_time": "2022-11-20T01:02:29.164918",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ids = []\n",
    "for id_ in test_ds['id']:\n",
    "    for col in LABEL_COLS:\n",
    "        ids.append(f\"{id_}_{col}\")\n",
    "\n",
    "predictions = post_process_logits(fin_logits, threshold=0.5)\n",
    "print(predictions.shape)\n",
    "sub_df['id'] = ids\n",
    "sub_df['predictions'] = predictions\n",
    "sub_df.to_csv(SUB_OUT_DIR/f\"{cfg.exp_name}_sub.csv\", index=False)\n",
    "sub_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Upload artifacts to W&B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_artifact = wandb.Artifact(name=cfg.exp_name, type=\"model\")\n",
    "infer_artifact.add_dir(SUB_OUT_DIR)\n",
    "wandb.log_artifact(infer_artifact, aliases=\"submission\")\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.015329,
     "end_time": "2022-11-20T01:02:47.388235",
     "exception": false,
     "start_time": "2022-11-20T01:02:47.372906",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Clean-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 0.014869,
     "end_time": "2022-11-20T01:02:47.418166",
     "exception": false,
     "start_time": "2022-11-20T01:02:47.403297",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if ENVIRON == \"kaggle\":\n",
    "    shutil.rmtree(\"./tmp\", ignore_errors=True)\n",
    "    shutil.rmtree(ROOT_DIR, ignore_errors=True)"
   ]
  }
 ],
 "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.8.10"
  },
  "name": "kaggle.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}