train_vlsp_fastai.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Import lib, modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from datasets.mrc_ner_dataset_vlsp_fastai import SentenceTransform, YTransform, BeforeBatchTransform\n",
    "from utils.read_file import load_jsonl\n",
    "from transformers import PhobertTokenizer\n",
    "from fastai.text.all import (\n",
    "    Datasets,\n",
    "    tensor,\n",
    "    Learner,\n",
    "    Callback,\n",
    "    EarlyStoppingCallback, SaveModelCallback, ReduceLROnPlateau,\n",
    "    GradientAccumulation,\n",
    "    NativeMixedPrecision, store_attr,\n",
    "    # GradientClip,\n",
    ")\n",
    "import os\n",
    "from models.phobert_query_ner import PhoBertQueryNER, splitter\n",
    "from models.query_ner_config import PhobertQueryNerConfig\n",
    "from loss.loss_func_fastai import CustomLoss\n",
    "from metrics.query_span_f1_fastai import QuerySpanF1"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Config"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "# Config\n",
    "bert_model = 'vinai/phobert-base'\n",
    "bert_dropout = 0.1\n",
    "mrc_dropout = 0.3\n",
    "loss_type = 'bce'   # bce, dice\n",
    "flat_ner = False\n",
    "span_loss_candidates = 'pred_and_gold' # all, pred_and_gold, gold\n",
    "max_length = 210\n",
    "batch_size = 2\n",
    "weight_decay = 0.01\n",
    "weight_start = 1.0\n",
    "weight_end = 1.0\n",
    "weight_span = 0.1\n",
    "data_dir = 'data/vlsp_processed'"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Data"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load data"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "train_data = load_jsonl(os.path.join(data_dir, 'train.jsonl'))\n",
    "valid_data = load_jsonl(os.path.join(data_dir, 'dev.jsonl'))\n",
    "test_data = load_jsonl(os.path.join(data_dir, 'test.jsonl'))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Data Pipeline"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "train_valid_data = [*train_data, *valid_data]\n",
    "n_samples = len(train_valid_data)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "splits = (list(range(len(train_data))), range(len(train_data), n_samples))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embedding are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = PhobertTokenizer.from_pretrained(bert_model)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "sent_tfms = SentenceTransform(\n",
    "    tokenizer=tokenizer,\n",
    "    possible_only=False,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data = [sent_tfms(i) for i in train_valid_data]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "test_data = [sent_tfms(i) for i in test_data]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Save data"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [],
   "source": [
    "# torch.save(data, 'train_valid_data.pkl')\n",
    "# torch.save(test_data, 'test_data.pkl')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "data = torch.load('train_valid_data.pkl')\n",
    "test_data = torch.load('test_data.pkl')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Dataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [],
   "source": [
    "x_tfms = [lambda x: x['token_ids'], tensor]\n",
    "y_tfms = [YTransform()]\n",
    "ds = Datasets(\n",
    "    data,\n",
    "    splits=splits,\n",
    "    tfms=[x_tfms, y_tfms],\n",
    "    n_inp=1,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "data": {
      "text/plain": "((tensor([    0, 14110,    18,   646,   221,     4,   221,  6143,     6,    86,\n              7,    18,     4,   221,  2044,     4,  1116, 18395,     4, 38207,\n              2,  2522,   985,    26,  2796,     2]),\n  (tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n           0, 0]),\n   tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n           0, 0]),\n   tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,\n           1, 0]),\n   tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,\n           1, 0]),\n   tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0],\n           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n            0, 0]]))),\n None)"
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds[0], ds.show(ds[0])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Dataloader"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_dls(ds, bs):\n",
    "    return ds.dataloaders(\n",
    "        bs=bs,\n",
    "        before_batch=BeforeBatchTransform(\n",
    "            max_seq_length=max_length,\n",
    "            pad_fields=[0, 1, 2, 3, 4],\n",
    "            pad_values=[tokenizer.pad_token_id, 0, 0, 0, 0],\n",
    "            pad_token_id=tokenizer.pad_token_id,\n",
    "            sep_token_id=tokenizer.sep_token_id,\n",
    "        ),\n",
    "        n_inp=1,\n",
    "        verbose=5,\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting up after_item: Pipeline: \n",
      "Setting up before_batch: Pipeline: BeforeBatchTransform\n",
      "Setting up after_batch: Pipeline: \n"
     ]
    }
   ],
   "source": [
    "dls = get_dls(ds, batch_size)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [
    {
     "data": {
      "text/plain": "(tensor([[    0, 14110,   646,  5267,     4,  2877,     4,  6146,     4,  1674,\n              4,   110,   201,     2,  1615, 16491,     4,  8094, 37513,     4,\n           2124,   736,   621,     4,  1601,   125,   367,    81,   639,   247,\n             14,    30,   342,   247,    35,   367,    98,  5291,     5,     2],\n         [    0, 14110,   646,  5267,     4,  2877,     4,  6146,     4,  1674,\n              4,   110,   201,     2,  1351,   174,     4,  1755,    74,    50,\n             10,   204,    36,  1205,    24,    86, 30716, 38850,     4,    12,\n           1445,   424,   332,   734,  1567,     5,     2,     1,     1,     1]],\n        device='cuda:0'),\n tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], device='cuda:0'),\n tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], device='cuda:0'),\n tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],\n         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]], device='cuda:0'),\n tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],\n         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]], device='cuda:0'),\n tensor([[[0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0],\n          ...,\n          [0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0]],\n \n         [[0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0],\n          ...,\n          [0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0],\n          [0, 0, 0,  ..., 0, 0, 0]]], device='cuda:0'))"
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dls.one_batch()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Model"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class WrapperPhoBertQueryNER(PhoBertQueryNER):\n",
    "    def forward(self, tokens):\n",
    "        attention_mask = (tokens != self.pad_token_id).long()\n",
    "        return super().forward(input_ids=tokens, token_type_ids=None, attention_mask=attention_mask)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Callback"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class GradientClip(Callback):\n",
    "    \"Clip norm of gradients\"\n",
    "    order=NativeMixedPrecision.order+1\n",
    "    def __init__(self,max_norm:float=1., norm_type:float=2.0): store_attr()\n",
    "    def before_step(self): nn.utils.clip_grad_norm_(self.parameters(), self.max_norm, self.norm_type)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Training"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model_config = PhobertQueryNerConfig.from_pretrained(\n",
    "    bert_model,\n",
    "    hidden_dropouy_prob=bert_dropout,\n",
    "    attention_probs_dropout_prob=bert_dropout,\n",
    "    mrc_dropout=mrc_dropout,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at vinai/phobert-base were not used when initializing WrapperPhoBertQueryNER: ['lm_head.bias', 'lm_head.dense.weight', 'lm_head.dense.bias', 'lm_head.layer_norm.weight', 'lm_head.layer_norm.bias', 'lm_head.decoder.weight', 'lm_head.decoder.bias']\n",
      "- This IS expected if you are initializing WrapperPhoBertQueryNER 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 WrapperPhoBertQueryNER 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 WrapperPhoBertQueryNER were not initialized from the model checkpoint at vinai/phobert-base and are newly initialized: ['start_outputs.weight', 'start_outputs.bias', 'end_outputs.weight', 'end_outputs.bias', 'span_embedding.classifier1.weight', 'span_embedding.classifier1.bias', 'span_embedding.classifier2.weight', 'span_embedding.classifier2.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": [
    "model = WrapperPhoBertQueryNER.from_pretrained(bert_model, config=model_config)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "outputs": [],
   "source": [
    "loss_func = CustomLoss(\n",
    "    weight_start=weight_start,\n",
    "    weight_end=weight_end,\n",
    "    weight_span=weight_span,\n",
    "    span_loss_candidates=span_loss_candidates,\n",
    "    loss_type=loss_type,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "outputs": [],
   "source": [
    "learner = Learner(\n",
    "    dls=dls,\n",
    "    model=model,\n",
    "    loss_func=loss_func,\n",
    "    splitter=splitter,\n",
    "    metrics=[QuerySpanF1(flat=flat_ner),],\n",
    "    cbs=[\n",
    "        EarlyStoppingCallback(\n",
    "            monitor='query_span_f1',\n",
    "            min_delta=1e-4,\n",
    "            patience=10,\n",
    "        ),\n",
    "        SaveModelCallback(\n",
    "            monitor='query_span_f1',\n",
    "            min_delta=1e-4,\n",
    "            fname='best_query_span_f1_model',\n",
    "            with_opt=True,\n",
    "        ),\n",
    "        ReduceLROnPlateau(monitor='valid_loss', min_delta=1e-4, patience=2),\n",
    "        GradientAccumulation(n_acc=32),\n",
    "        GradientClip(max_norm=1., norm_type=2.0,)\n",
    "    ],\n",
    "    wd=weight_decay,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Freeze Roberta"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "learner.freeze_to(-1)\n",
    "learner.lr_find()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "learner.fit_one_cycle(2, 1e-3)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Unfreeze"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "learner.unfreeze()\n",
    "learner.dls = get_dls(ds, bs=8,)\n",
    "learner.lr_find()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "lr = 1e-4\n",
    "learner.fit_one_cycle(20, slice(lr / (2.6 **4), lr))\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}