train_pii_longtransformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "class LSTMHead(nn.Module):\n",
    "    def __init__(self, in_features, hidden_dim, n_layers):\n",
    "        super().__init__()\n",
    "        self.lstm = nn.LSTM(in_features,\n",
    "                            hidden_dim,\n",
    "                            n_layers,\n",
    "                            batch_first=True,\n",
    "                            bidirectional=True,\n",
    "                            dropout=0.)\n",
    "        self.out_features = hidden_dim\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.lstm.flatten_parameters()\n",
    "        hidden, (_, _) = self.lstm(x)\n",
    "        out = hidden\n",
    "        return out\n",
    "\n",
    "\n",
    "class PIIModel(pl.LightningModule):\n",
    "    def __init__(self,config):\n",
    "        super().__init__()\n",
    "        self.cfg = config\n",
    "        self.model_config = AutoConfig.from_pretrained(\n",
    "            config.model_name,\n",
    "        )\n",
    "        hidden_dropout_prob: float = 0.1\n",
    "        layer_norm_eps: float = 1e-7\n",
    "        self.model_config.update(\n",
    "            {\n",
    "                \"output_hidden_states\": True,\n",
    "                \"hidden_dropout_prob\": hidden_dropout_prob,\n",
    "                \"layer_norm_eps\": layer_norm_eps,\n",
    "                \"add_pooling_layer\": False,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        self.transformers_model = AutoModel.from_pretrained(config.model_name,config=self.model_config)\n",
    "        self.head = LSTMHead(in_features=self.model_config.hidden_size, hidden_dim=self.model_config.hidden_size//2, n_layers=1)\n",
    "\n",
    "        self.output = nn.Linear(self.model_config.hidden_size, len(self.cfg.target_cols))\n",
    "\n",
    "        self.loss_function = nn.CrossEntropyLoss(reduction='mean',ignore_index=-100) \n",
    "        self.validation_step_outputs = []\n",
    "        \n",
    "\n",
    "\n",
    "    def forward(self, input_ids, attention_mask,train):\n",
    "        \n",
    "        transformer_out = self.transformers_model(input_ids,attention_mask = attention_mask)\n",
    "        sequence_output = transformer_out.last_hidden_state\n",
    "        sequence_output = self.head(sequence_output)\n",
    "        logits = self.output(sequence_output)\n",
    "        return (logits, _)\n",
    "\n",
    "    def training_step(self,batch,batch_idx):\n",
    "        input_ids = batch['input_ids']\n",
    "        attention_mask = batch['attention_mask']\n",
    "        target = batch['labels'] \n",
    "\n",
    "        outputs = self(input_ids,attention_mask,train=True)\n",
    "        output = outputs[0]\n",
    "\n",
    "        # loss = self.loss_function(output.view(-1, len(self.cfg.target_cols)),target.view(-1))\n",
    "        loss = self.loss_function(output.view(-1,len(self.cfg.target_cols)), target.view(-1))\n",
    "\n",
    "        self.log('train_loss', loss , prog_bar=True)\n",
    "        return {'loss': loss}\n",
    "    \n",
    "    def train_epoch_end(self,outputs):\n",
    "        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()\n",
    "        print(f'epoch {trainer.current_epoch} training loss {avg_loss}')\n",
    "        return {'train_loss': avg_loss} \n",
    "    \n",
    "    def train_dataloader(self):\n",
    "        return self._train_dataloader \n",
    "    \n",
    "    def validation_dataloader(self):\n",
    "        return self._validation_dataloader\n",
    "\n",
    "    def get_optimizer_params(self, encoder_lr, decoder_lr, weight_decay=0.0):\n",
    "        param_optimizer = list(model.named_parameters())\n",
    "        no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
    "        optimizer_parameters = [\n",
    "            {'params': [p for n, p in self.transformers_model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
    "             'lr': encoder_lr, 'weight_decay': weight_decay},\n",
    "            {'params': [p for n, p in self.transformers_model.named_parameters() if any(nd in n for nd in no_decay)],\n",
    "             'lr': encoder_lr, 'weight_decay': 0.0},\n",
    "            {'params': [p for n, p in self.named_parameters() if \"transformers_model\" not in n],\n",
    "             'lr': decoder_lr, 'weight_decay': 0.0}\n",
    "        ]\n",
    "        return optimizer_parameters\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = AdamW(self.parameters(), lr = config.learning_rate)\n",
    "\n",
    "        epoch_steps = self.cfg.data_length\n",
    "        batch_size = self.cfg.batch_size\n",
    "\n",
    "        warmup_steps = 0.1 * epoch_steps // batch_size\n",
    "        training_steps = self.cfg.epochs * epoch_steps // batch_size\n",
    "        # scheduler = get_linear_schedule_with_warmup(optimizer,warmup_steps,training_steps,-1)\n",
    "        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer, warmup_steps, training_steps, lr_end=1e-6, power=3.0)\n",
    "\n",
    "        lr_scheduler_config = {\n",
    "                'scheduler': scheduler,\n",
    "                'interval': 'step',\n",
    "                'frequency': 1,\n",
    "            }\n",
    "\n",
    "        return {'optimizer': optimizer, 'lr_scheduler': lr_scheduler_config}"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}