train_gpt2_summarizer_np.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from datetime import datetime\n",
    "import os\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "from transformers import GPT2LMHeadModel,AdamW, WarmupLinearSchedule\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import torch\n",
    "from torch.nn import CrossEntropyLoss\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
    "from tqdm import tnrange, tqdm_notebook\n",
    "\n",
    "from dataset import GPT21024Dataset \n",
    "from utils import add_special_tokens, generate_sample, set_seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#please change default arguments if needed\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--lr\",default=5e-5, type=float, help=\"learning rate\")\n",
    "parser.add_argument(\"--seed\",default=42, type=int,  help=\"seed to replicate results\")\n",
    "parser.add_argument(\"--n_gpu\",default=1, type=int,  help=\"no of gpu available\")\n",
    "parser.add_argument(\"--gradient_accumulation_steps\",default=2, type=int, help=\"gradient_accumulation_steps\")\n",
    "parser.add_argument(\"--batch_size\",default=1, type=int,  help=\"batch_size\")\n",
    "parser.add_argument(\"--num_workers\",default=4, type=int,  help=\"num of cpus available\")\n",
    "parser.add_argument(\"--device\",default=torch.device('cpu'), help=\"torch.device object\")\n",
    "parser.add_argument(\"--num_train_epochs\",default=1, type=int,  help=\"no of epochs of training\")\n",
    "parser.add_argument(\"--output_dir\",default='./output', type=str,  help=\"path to save evaluation results\")\n",
    "parser.add_argument(\"--model_dir\",default='./weights', type=str,  help=\"path to save trained model\")\n",
    "parser.add_argument(\"--max_grad_norm\",default=1.0, type=float, help=\"max gradient norm.\")\n",
    "parser.add_argument(\"--root_dir\",default='./CNN/gpt2_1024_data', type=str, help=\"location of json dataset.\")\n",
    "parser.add_argument(\"--ids_file\",default='./CNN/ids.json', type=str, help=\"location of train, valid and test file indexes\")\n",
    "args = parser.parse_args([])\n",
    "print(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(args, tokenizer, model, data, ignore_index, train_size=3000, val_size=500):\n",
    "\t\"\"\" Trains GPT2 model and logs necessary details.\n",
    "\t\tArgs:\n",
    "\t\t\targs: dict that contains all the necessary information passed by user while training\n",
    " \t\t\tmodel: finetuned gpt/gpt2 model\n",
    "\t\t\ttokenizer: GPT/GPT2 tokenizer\n",
    "\t\t\tdata: numpy_data created using prepare_np_data.py\n",
    "\t\t\tignore_index: token not considered in loss calculation\n",
    "            train_size = size of training data\n",
    "            val_size = size of validation data\n",
    "\t\"\"\"\n",
    "    writer = SummaryWriter('./output/logs')\n",
    "    train_dataset = TensorDataset(torch.from_numpy(data[:train_size]).type('torch.LongTensor'))\n",
    "    valid_dataset = TensorDataset(torch.from_numpy(data[train_size:train_size+val_size]).type('torch.LongTensor'))\n",
    "    train_sampler = RandomSampler(train_dataset)\n",
    "    train_dl = DataLoader(train_dataset,sampler=train_sampler,batch_size=args.batch_size,num_workers=args.num_workers)\n",
    "    loss_fct = CrossEntropyLoss(ignore_index=ignore_index) #ignores padding token for loss calculation\n",
    "    optimizer = AdamW(model.parameters(),lr=args.lr)\n",
    "    scheduler = WarmupLinearSchedule(optimizer,100,80000)\n",
    "\n",
    "    global_step = 0\n",
    "    tr_loss, logging_loss = 0.0, 0.0\n",
    "    model.zero_grad()\n",
    "    train_iterator = tnrange(int(args.num_train_epochs), desc=\"Epoch\")\n",
    "    set_seed(args)\n",
    "    for _ in train_iterator:\n",
    "        epoch_iterator = tqdm_notebook(train_dl, desc=\"Training\")\n",
    "        for step, batch in enumerate(epoch_iterator):\n",
    "            inputs, labels = batch['article'].to(args.device), batch['article'].to(args.device)\n",
    "            model.train()\n",
    "            logits = model(inputs)[0]\n",
    "            idx = batch['sum_idx'].item() # index of separator token\n",
    "            # only consider loss on reference summary just like seq2seq models\n",
    "            shift_logits = logits[..., idx:-1, :].contiguous()\n",
    "            shift_labels = labels[..., idx+1:].contiguous()\n",
    "            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n",
    "            loss = loss/args.gradient_accumulation_steps\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)\n",
    "            tr_loss += loss.item()\n",
    "            if (step + 1) % args.gradient_accumulation_steps == 0:\n",
    "                optimizer.step()\n",
    "                scheduler.step()  # Update learning rate schedule\n",
    "                model.zero_grad()\n",
    "                global_step += 1\n",
    "                writer.add_scalar('lr', scheduler.get_lr()[0], global_step)\n",
    "                writer.add_scalar('loss', (tr_loss - logging_loss)/args.gradient_accumulation_steps, global_step)\n",
    "                logging_loss = tr_loss\n",
    "                print(\"loss:\", loss.item(), end='\\n\\n')\n",
    "                if (step + 1)/args.gradient_accumulation_steps == 1.0:\n",
    "                \tprint('After 1st update: ', end='\\n\\n')\n",
    "                \tgenerate_sample(valid_dataset, tokenizer, num=2, eval_step=False,devcie=args.device)\n",
    "                \n",
    "                \n",
    "            if (step + 1) % (10*args.gradient_accumulation_steps) == 0:\n",
    "                results = evaluate(args, model, valid_dataset, ignore_index, global_step)\n",
    "                for key, value in results.items():\n",
    "                    writer.add_scalar('eval_{}'.format(key), value, global_step)\n",
    "                print('After', global_step+1,'updates: ', end='\\n\\n')\n",
    "                generate_sample(valid_dataset, tokenizer, num=2, eval_step=True,device=args.device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(args, model, eval_dataset, ignore_index, global_step=None):\n",
    "    \"\"\" Returns perplexity score on validation dataset.\n",
    "        Args:\n",
    "            args: dict that contains all the necessary information passed by user while training\n",
    "            model: finetuned gpt/gpt2 model\n",
    "            eval_dataset: GPT21024Dataset object for validation data\n",
    "            global_step: no. of times gradients have backpropagated\n",
    "            ignore_index: token not considered in loss calculation\n",
    "    \"\"\"\n",
    "    if not os.path.exists(args.output_dir):\n",
    "        os.mkdir(args.output_dir)\n",
    "    eval_output_dir = args.output_dir\n",
    "\n",
    "    results = {}\n",
    "    eval_sampler = SequentialSampler(eval_dataset)\n",
    "    eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.batch_size)\n",
    "    loss_fct = CrossEntropyLoss(ignore_index=ignore_index) #ignores padding token for loss calculation\n",
    "\n",
    "    eval_loss = 0.0\n",
    "    nb_eval_steps = 0\n",
    "    model.eval()\n",
    "\n",
    "    for batch in tqdm_notebook(eval_dataloader, desc=\"Evaluating\"):\n",
    "        inputs, labels = torch.tensor(batch['article']).to(args.device), torch.tensor(batch['article']).to(args.device)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "        \tlogits = model(inputs)[0]\n",
    "            idx = batch['sum_idx'].item() # index of separator token\n",
    "            # only consider loss on reference summary just like seq2seq models\n",
    "            shift_logits = logits[..., idx:-1, :].contiguous()\n",
    "            shift_labels = labels[..., idx+1:].contiguous()\n",
    "            lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n",
    "            eval_loss += lm_loss.mean().item()\n",
    "        nb_eval_steps += 1\n",
    "\n",
    "    eval_loss = eval_loss / nb_eval_steps\n",
    "    perplexity = torch.exp(torch.tensor(eval_loss))\n",
    "\n",
    "    result = {\n",
    "        \"perplexity\": perplexity\n",
    "    }\n",
    "    print(\"perplexity:\", perplexity.item())\n",
    "\n",
    "    if global_step:\n",
    "    \toutput_eval_file = os.path.join(eval_output_dir, \"eval_results.txt\")\n",
    "\t    with open(output_eval_file, \"a\") as f:\n",
    "\t        for key in sorted(result.keys()):\n",
    "\t            f.write('\\n\\n')\n",
    "\t            f.write(\"time = %s, %s = %s, step = %s\\n\" % (datetime.now().strftime(\"%d/%m/%Y %H:%M:%S\"), key, str(result[key]), str(global_step)))\n",
    "    return result           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reading dataset\n",
    "data = np.load(file_name)\n",
    "train_size, val_size, test_size = int(len(data)*0.9), int(len(data)*0.05), int(len(data)*0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load pretrained GPT2\n",
    "tokenizer = add_special_tokens()\n",
    "ignore_idx = tokenizer.pad_token_id\n",
    "model = GPT2LMHeadModel.from_pretrained('gpt2')\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "model.to(args.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training\n",
    "\n",
    "start = time.time()\n",
    "train(args, tokenizer, model, data, ignore_index, train_size, val_size)\n",
    "print('total time: ', (time.time()-start)/60, \" minutes\", end='\\n\\n')\n",
    "\n",
    "print('Saving trained model...')\n",
    "model_file = os.path.join(args.model_dir, 'model_data{}_trained_after_{}_epochs_only_sum_loss_ignr_pad.bin'.format(train_size,args.num_train_epochs))\n",
    "config_file = os.path.join(args.model_dir, 'config_data{}_trained_after_{}_epochs_only_sum_loss_ignr_pad.json'.format(train_size,args.num_train_epochs))\n",
    "torch.save(model.state_dict(), model_file)\n",
    "model.config.to_json_file(config_file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}