firebert_imdb_and_adversarial_co-tuner.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coding=utf-8\n",
    "# Copyright 2020 FireBERT authors. All rights reserved.\n",
    "#\n",
    "# Licensed under the MIT license\n",
    "# See https://github.com/FireBERT-author/FireBERT/blob/master/LICENSE for details\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "#"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from processors import ImdbProcessor\n",
    "from firebert_fct import FireBERT_FCT\n",
    "from bert_base_model import LightningBertForSequenceClassification\n",
    "\n",
    "num_gpus = -1 if torch.cuda.is_available() else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# place where we will save our model\n",
    "\n",
    "save_root_path ='resources/models/co-tuned_IMDB_on_lightning/'\n",
    "\n",
    "use_full_example = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare hyperparameters\n",
    "\n",
    "max_steps = -1 # if -1 then calculate number of training steps based on the length of the train set\n",
    "len_train_set = 40000\n",
    "\n",
    "gradient_accumulation_steps = 1\n",
    "learning_rate = 2e-5\n",
    "weight_decay = 0.00\n",
    "adam_epsilon = 1e-8\n",
    "warmup_proportion = 0 \n",
    "\n",
    "num_train_epochs = 1\n",
    "batch_size = 2\n",
    "\n",
    "if max_steps > 0:\n",
    "    num_train_epochs = max_steps // (len_train_set // gradient_accumulation_steps) + 1\n",
    "    num_training_steps = max_steps\n",
    "else:\n",
    "    num_training_steps = len_train_set // gradient_accumulation_steps * num_train_epochs\n",
    "    \n",
    "warmup_steps = num_training_steps // num_train_epochs * warmup_proportion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "40000"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_training_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "warmup_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters for generating adversarial candidates during co-tuning\n",
    "\n",
    "use_USE = True\n",
    "USE_method = 'filter'\n",
    "USE_multiplier = 12\n",
    "stop_words = True\n",
    "perturb_words = 23\n",
    "candidates_per_word = 10\n",
    "total_alternatives = 4\n",
    "match_pos = True\n",
    "leave_alone = 0\n",
    "random_out_of = 0\n",
    "judge_bert = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:absl:Using scratch/tf_cache to cache modules.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n"
     ]
    }
   ],
   "source": [
    "hparams = { 'learning_rate': learning_rate,\n",
    "            'adam_epsilon': adam_epsilon,\n",
    "            'weight_decay': weight_decay,\n",
    "            'warmup_steps': warmup_steps,\n",
    "            'num_training_steps': num_training_steps,\n",
    "            'batch_size': batch_size,\n",
    "            'use_USE': use_USE,\n",
    "            'USE_method': USE_method,\n",
    "            'USE_multiplier': USE_multiplier,\n",
    "            'stop_words': stop_words,\n",
    "            'perturb_words': perturb_words,\n",
    "            'candidates_per_word': candidates_per_word,\n",
    "            'total_alternatives': total_alternatives,\n",
    "            'match_pos': match_pos,\n",
    "            'use_full_example': use_full_example,\n",
    "            'leave_alone': leave_alone,\n",
    "            'random_out_of': random_out_of,\n",
    "            'judge_bert': judge_bert\n",
    "           }\n",
    "\n",
    "proc_hparams = {}\n",
    "# delete this next line to run full 100%\n",
    "#proc_hparams.update({'sample_percent': 3,\n",
    "#                     'randomize': False})\n",
    "\n",
    "# instantiate the model used for SWITCH\n",
    "switch_model = LightningBertForSequenceClassification(load_from = 'resources/models/IMDB/pytorch_model.bin', \n",
    "                                                      processor = ImdbProcessor(), \n",
    "                                                      hparams = {'batch_size': 6 })\n",
    "\n",
    "switch_model.cuda()\n",
    "#switch_model = None\n",
    "\n",
    "\n",
    "model = FireBERT_FCT(switch_model=switch_model,\n",
    "                     processor=ImdbProcessor(hparams=proc_hparams), hparams=hparams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = model.get_processor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset, train_examples = \\\n",
    "        processor.load_and_cache_examples(\"data/IMDB\", example_set='train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_dataset, _ = processor.load_and_cache_examples(\"data/IMDB\", example_set='dev')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset, _ = processor.load_and_cache_examples(\"data/IMDB\", example_set='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_train_dataset(train_dataset, train_examples)\n",
    "model.set_val_dataset(val_dataset)\n",
    "model.set_test_dataset(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning.logging import TensorBoardLogger\n",
    "\n",
    "tensor_logger = TensorBoardLogger(save_dir= save_root_path + 'logs', version=10, name='imdb_finetuning')\n",
    "checkpoint_save_path = save_root_path + 'checkpoints/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    filepath=checkpoint_save_path,\n",
    "    verbose=True,\n",
    "    monitor='val_loss',\n",
    "    mode='min'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "amp_opt_level='O1' # https://nvidia.github.io/apex/amp.html#opt-levels\n",
    "max_grad_norm = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(default_save_path=checkpoint_save_path, logger=tensor_logger, gpus=num_gpus,\n",
    "                     max_epochs = num_train_epochs, amp_level=amp_opt_level, gradient_clip_val=max_grad_norm,\n",
    "                     max_steps = num_training_steps, checkpoint_callback=checkpoint_callback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validation sanity check', layout=Layout(flex='2'), max=5.…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ccd20295c764b2c9670343f403178b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), max=1.0), HTML(value='')), …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=2500.0, style=P…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'val_loss': 0.2741829454898834, 'avg_val_acc': 0.8970000147819519}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.tqdm_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('resources/models/co-tuned_IMDB_on_lightning_final/vocab.txt',\n",
       " 'resources/models/co-tuned_IMDB_on_lightning_final/special_tokens_map.json',\n",
       " 'resources/models/co-tuned_IMDB_on_lightning_final/added_tokens.json')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.save_checkpoint(save_root_path + 'training_checkpoint')\n",
    "\n",
    "torch.save(model.state_dict(), save_root_path + 'pytorch_model.bin')\n",
    "with open(save_root_path + 'bert_config.json', 'w') as f:\n",
    "    f.write(model.bert.config.to_json_string())\n",
    "model.tokenizer.save_pretrained(save_root_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30301bf833024209a7045da3fbe784b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Testing', layout=Layout(flex='2'), max=2500.0, style=Prog…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------\n",
      "TEST RESULTS\n",
      "{'avg_test_acc': tensor(0.8984)}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "trainer.test(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'val_loss': 0.2741829454898834,\n",
       " 'avg_val_acc': 0.8970000147819519,\n",
       " 'avg_test_acc': 0.8984000086784363}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.tqdm_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset\n",
    "\n",
    "def load_examples(features_file):\n",
    "\n",
    "    features = torch.load(features_file)\n",
    "\n",
    "    # Convert to Tensors and build dataset\n",
    "    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)\n",
    "    all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)\n",
    "    all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)\n",
    "    all_labels = torch.tensor([f.label for f in features], dtype=torch.long)\n",
    "\n",
    "    all_idxs = torch.tensor([i for i in range(len(all_input_ids))], dtype=torch.long)\n",
    "        \n",
    "    dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels, all_idxs)\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare how well the model does against adversarial samples\n",
    "test_dataset = load_examples('data/IMDB/generated/imdb_adversarial_samples_for_dev')\n",
    "model.set_test_dataset(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "231a5b8cd84348e8955bebafe1a5b497",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Testing', layout=Layout(flex='2'), max=2071.0, style=Prog…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------\n",
      "TEST RESULTS\n",
      "{'avg_test_acc': tensor(0.8728)}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'avg_test_acc': 0.8727667927742004}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = pl.Trainer(gpus=num_gpus)\n",
    "trainer.test(model)\n",
    "trainer.tqdm_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0df9eeb97034da086223824e78a760a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Testing', layout=Layout(flex='2'), max=2051.0, style=Prog…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------\n",
      "TEST RESULTS\n",
      "{'avg_test_acc': tensor(0.8776)}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'avg_test_acc': 0.8776206970214844}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# compare how well the model does against test adversarial samples\n",
    "test_dataset = load_examples('data/IMDB/generated/imdb_adversarial_samples_for_test')\n",
    "model.set_test_dataset(test_dataset)\n",
    "trainer = pl.Trainer(gpus=num_gpus)\n",
    "trainer.test(model)\n",
    "trainer.tqdm_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "438527f22c4b4877bcbafa77a9605d48",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Testing', layout=Layout(flex='2'), max=2500.0, style=Prog…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------\n",
      "TEST RESULTS\n",
      "{'avg_test_acc': tensor(0.8970)}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'avg_test_acc': 0.8970000147819519}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load up the dev samples again to get eval timings (that we didn't get from training)\n",
    "test_dataset = load_examples('data/IMDB/cached_dev_bert-base-uncased_256_imdb')\n",
    "model.set_test_dataset(test_dataset)\n",
    "trainer = pl.Trainer(gpus=num_gpus)\n",
    "trainer.test(model)\n",
    "trainer.tqdm_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}