seqganbert-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "daily-listing",
   "metadata": {},
   "source": [
    "### Import all necessary packages "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7383caf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import io\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "import numpy as np\n",
    "import time\n",
    "import math\n",
    "import datetime\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n",
    "import os\n",
    "from transformers import AutoModel,AutoTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "outdoor-warren",
   "metadata": {},
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "studied-temperature",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_qc_examples(input_file, num_sentences=None):\n",
    "    examples = []\n",
    "\n",
    "    with open(input_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "        if(num_sentences==0):\n",
    "            return examples\n",
    "        if num_sentences:\n",
    "            lines = lines[:num_sentences]\n",
    "\n",
    "        for line in lines:\n",
    "            split = line.split(\"\\t\")\n",
    "            text = split[0]\n",
    "            labels = split[1][:-1]\n",
    "            examples.append((text, labels))\n",
    "\n",
    "    return examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "personalized-castle",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data_loader(input_examples, label_masks, label_map, do_shuffle = False, balance_label_examples = False):\n",
    "    '''\n",
    "    Generate a Dataloader given the input examples, eventually masked if they are to be considered NOT labeled.\n",
    "    '''\n",
    "    examples = []\n",
    "\n",
    "    # Count the percentage of labeled examples\n",
    "    num_labeled_examples = 0\n",
    "    for label_mask in label_masks:\n",
    "        if label_mask:\n",
    "            num_labeled_examples += 1\n",
    "\n",
    "    label_mask_rate = num_labeled_examples/len(input_examples)\n",
    "\n",
    "    # if required it applies the balance\n",
    "    for index, ex in enumerate(input_examples):\n",
    "        if label_mask_rate == 1 or not balance_label_examples:\n",
    "            examples.append((ex, label_masks[index]))\n",
    "        else:\n",
    "            # IT SIMULATE A LABELED EXAMPLE\n",
    "            if label_masks[index]:\n",
    "                balance = int(1/label_mask_rate)\n",
    "                balance = int(math.log(balance,2))\n",
    "                if balance < 1:\n",
    "                    balance = 1\n",
    "                for b in range(0, int(balance)):\n",
    "                    examples.append((ex, label_masks[index]))\n",
    "            else:\n",
    "                examples.append((ex, label_masks[index]))\n",
    "\n",
    "    input_ids = []\n",
    "    input_mask_array = []\n",
    "    label_mask_array = []\n",
    "    label_id_array = []\n",
    "    label_id01_array = []\n",
    "\n",
    "    # Tokenization\n",
    "    for (text, label_mask) in examples:\n",
    "        # text[0] is the sentence\n",
    "        # text[1] is the list of 1s and 0s separated by space mentioning which word is a disfluency\n",
    "        # label_mask is True (if sentence is labelled) or False\n",
    "        \n",
    "        encoded_sent = tokenizer.encode(text[0], add_special_tokens=True, max_length=max_seq_length, padding=\"max_length\", truncation=True)\n",
    "        input_ids.append(encoded_sent)\n",
    "        label_mask_array.append(label_mask)\n",
    "\n",
    "        tokenized_inputs = tokenizer([text[0].split()], max_length=max_seq_length, padding=\"max_length\", truncation=True, is_split_into_words=True)\n",
    "        word_ids = tokenized_inputs.word_ids(batch_index=0)\n",
    "\n",
    "        previous_word_idx = None\n",
    "        label_ids = []\n",
    "        for word_idx in word_ids:\n",
    "            # Special tokens have a word id that is None. We set the label to -100 so they are automatically\n",
    "            # ignored in the loss function.\n",
    "            if word_idx is None:\n",
    "                label_ids.append(PAD_VALUE)\n",
    "            # We set the label for the first token of each word.\n",
    "            elif word_idx != previous_word_idx:\n",
    "                \n",
    "                if label_mask == 1:\n",
    "                    label_ids.append(int(text[1].split()[word_idx]))\n",
    "                else:\n",
    "                    label_ids.append(int(0)) # anything\n",
    "            # For the other tokens in a word, we set the label to either the current label or -100, depending on\n",
    "            # the label_all_tokens flag.\n",
    "            else:\n",
    "                if label_mask == 1:\n",
    "                    label_ids.append(int(text[1].split()[word_idx]))\n",
    "                else:\n",
    "                    label_ids.append(int(0)) # anything\n",
    "\n",
    "            previous_word_idx = word_idx\n",
    "\n",
    "        label_ids01 = [0 if id == PAD_VALUE else id for id in label_ids]\n",
    "\n",
    "        label_id_array.append(label_ids)\n",
    "        label_id01_array.append(label_ids01)\n",
    "\n",
    "\n",
    "    # Attention to token (to ignore padded input wordpieces)\n",
    "    for sent in input_ids:\n",
    "        att_mask = [int(token_id > 0) for token_id in sent]\n",
    "        input_mask_array.append(att_mask)\n",
    "    \n",
    "\n",
    "    input_ids = torch.tensor(input_ids) # List of list. Inner list contains ids for every word of the sentence and 0 for padding\n",
    "    input_mask_array = torch.tensor(input_mask_array) # List of list. Inner list contains 1 for every non zero input id word in example sentence\n",
    "    label_id_array = torch.tensor(label_id_array, dtype=torch.long) # 2 for padding, 1 for disfluent, 0 for non disfluent\n",
    "    label_id01_array = torch.tensor(label_id01_array, dtype=torch.long) # 1 for disfluent, 0 for padding/non disfluent\n",
    "    label_mask_array = torch.tensor(label_mask_array) # list of bool. True for sent with labeling, false for sent without labeling\n",
    "\n",
    "    dataset = TensorDataset(input_ids, input_mask_array, label_id_array, label_mask_array, label_id01_array)\n",
    "\n",
    "    if do_shuffle:\n",
    "        sampler = RandomSampler\n",
    "    else:\n",
    "        sampler = SequentialSampler\n",
    "\n",
    "    return DataLoader(\n",
    "              dataset,\n",
    "              sampler = sampler(dataset),\n",
    "              batch_size = batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "healthy-bumper",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_time(elapsed):\n",
    "    elapsed_rounded = int(round((elapsed)))\n",
    "    return str(datetime.timedelta(seconds=elapsed_rounded)) # Format as hh:mm:ss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cognitive-screening",
   "metadata": {},
   "source": [
    "### Define parameters for data usage and training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "vertical-solution",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed_val = 42\n",
    "random.seed(seed_val)\n",
    "np.random.seed(seed_val)\n",
    "torch.manual_seed(seed_val)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(seed_val)\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:3\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "    \n",
    "EXPERIMENT_ID= 'presto-english-small-230523-2'\n",
    "CHECKPOINT_DIR = f'checkpoints/{EXPERIMENT_ID}'\n",
    "transformer_model_name = \"google/muril-base-cased\" # Check Hugging Face for experimenting with other multilingual transformers\n",
    "\n",
    "\n",
    "label_list = [\"1\", \"0\"]\n",
    "PAD_VALUE = 2\n",
    "\n",
    "labeled_file = \"./data/sample_data/sample-train-labeled.tsv\"\n",
    "unlabeled_file = \"./data/sample_data/sample-train-unlabeled.tsv\"\n",
    "valid_filename = \"./data/sample_data/sample-valid-labeled.tsv\"\n",
    "test_filename = \"./data/sample_data/sample-test-labeled.tsv\"\n",
    "\n",
    "NUM_TRAINING_SENTENCES = len(open(labeled_file,'r').readlines()) # You can also specify a fixed number like 5000 or 10000\n",
    "NUM_UNLABELED_SENTENCES = len(open(unlabeled_file,'r').readlines()) # You can also specify a fixed number like 5000 or 10000\n",
    "\n",
    "print(\"Labeled file:\", labeled_file)\n",
    "print(\"Unlabeled file:\", unlabeled_file)\n",
    "print(\"Valid file:\", valid_filename)\n",
    "print(\"Test file:\", test_filename)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "annual-aquarium",
   "metadata": {},
   "source": [
    "### Define transformer parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cheap-parade",
   "metadata": {},
   "outputs": [],
   "source": [
    "#--------------------------------\n",
    "#  Transformer parameters\n",
    "#--------------------------------\n",
    "max_seq_length = 64\n",
    "batch_size = 16\n",
    "\n",
    "# number of hidden layers in the discriminator\n",
    "num_hidden_layers_d = 1; \n",
    "\n",
    "# size of the generator's input noisy vectors\n",
    "noise_size = 100\n",
    "\n",
    "# dropout to be applied to discriminator's input vectors\n",
    "out_dropout_rate = 0.4\n",
    "\n",
    "# Replicate labeled data to balance poorly represented datasets\n",
    "apply_balance = True\n",
    "\n",
    "#--------------------------------\n",
    "#  Optimization parameters\n",
    "#--------------------------------\n",
    "learning_rate_discriminator = 2e-5\n",
    "learning_rate_generator = 2e-5\n",
    "epsilon = 1e-8\n",
    "num_train_epochs = 40\n",
    "\n",
    "# Scheduler\n",
    "apply_scheduler = False\n",
    "warmup_proportion = 0.1\n",
    "\n",
    "# Print\n",
    "print_each_n_step = 10\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stainless-denmark",
   "metadata": {},
   "source": [
    "### Load transformer, tokenizer and data for training/testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f7c1f0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = AutoModel.from_pretrained(transformer_model_name)\n",
    "tokenizer = AutoTokenizer.from_pretrained(transformer_model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f0c518c",
   "metadata": {},
   "outputs": [],
   "source": [
    "labeled_examples = get_qc_examples(labeled_file, num_sentences=NUM_TRAINING_SENTENCES)\n",
    "unlabeled_examples = get_qc_examples(unlabeled_file, num_sentences=NUM_UNLABELED_SENTENCES)\n",
    "valid_examples = get_qc_examples(valid_filename)\n",
    "test_examples = get_qc_examples(test_filename)\n",
    "\n",
    "print(f\"Labeled Training Examples: {len(labeled_examples)}, Unlabeled Training Examples: {len(unlabeled_examples)}, Validation Examples: {len(valid_examples)}, Test Examples: {len(test_examples)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e98e031",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_map = {}\n",
    "for (i, label) in enumerate(label_list):\n",
    "    label_map[label] = i\n",
    "\n",
    "    \n",
    "train_examples = labeled_examples\n",
    "\n",
    "#The labeled (train) dataset is assigned with a mask set to True\n",
    "train_label_masks = np.ones(len(labeled_examples), dtype=bool)\n",
    "\n",
    "#If unlabel examples are available\n",
    "if unlabeled_examples:\n",
    "    train_examples = train_examples + unlabeled_examples\n",
    "    #The unlabeled (train) dataset is assigned with a mask set to False\n",
    "    tmp_masks = np.zeros(len(unlabeled_examples), dtype=bool)\n",
    "    train_label_masks = np.concatenate([train_label_masks,tmp_masks])\n",
    "\n",
    "train_dataloader = generate_data_loader(train_examples, train_label_masks, label_map, do_shuffle = True, balance_label_examples = apply_balance)\n",
    "\n",
    "valid_label_masks = np.ones(len(valid_examples), dtype=bool)\n",
    "valid_dataloader = generate_data_loader(valid_examples, valid_label_masks, label_map, do_shuffle = False, balance_label_examples = False)\n",
    "\n",
    "test_label_masks = np.ones(len(test_examples), dtype=bool)\n",
    "test_dataloader = generate_data_loader(test_examples, test_label_masks, label_map, do_shuffle = False, balance_label_examples = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "moral-vintage",
   "metadata": {},
   "source": [
    "### Define model based on hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84172aba",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, noise_size=100, hidden_size=512, dropout_rate=0.1):\n",
    "        super(Generator, self).__init__()\n",
    "        decoder_layer = nn.TransformerDecoderLayer(d_model=hidden_size, nhead=8)\n",
    "        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)\n",
    "\n",
    "    def forward(self, noise, memory):\n",
    "        return self.transformer_decoder(noise, memory)\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_size=512, hidden_sizes=[512], num_labels=2, dropout_rate=0.1):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.input_dropout = nn.Dropout(p=dropout_rate)\n",
    "        layers = []\n",
    "        hidden_sizes = [input_size] + hidden_sizes\n",
    "        for i in range(len(hidden_sizes)-1):\n",
    "            layers.extend([nn.Linear(hidden_sizes[i], hidden_sizes[i+1]), nn.LeakyReLU(0.2, inplace=True), nn.Dropout(dropout_rate)])\n",
    "\n",
    "        self.layers = nn.Sequential(*layers)\n",
    "        self.logit = nn.Linear(hidden_sizes[-1],num_labels+1) # +1 for the probability of this sample being fake/real.\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, input_rep):\n",
    "        input_rep = self.input_dropout(input_rep)\n",
    "        last_rep = self.layers(input_rep)\n",
    "        logits = self.logit(last_rep)\n",
    "        probs = self.softmax(logits)\n",
    "        return last_rep, logits, probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fede6073",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoConfig\n",
    "config = AutoConfig.from_pretrained(model_name)\n",
    "\n",
    "hidden_size = int(config.hidden_size)\n",
    "hidden_levels_d = [hidden_size for i in range(0, num_hidden_layers_d)]\n",
    "\n",
    "generator = Generator(noise_size=noise_size, hidden_size=hidden_size, dropout_rate=out_dropout_rate)\n",
    "discriminator = Discriminator(input_size=hidden_size, hidden_sizes=hidden_levels_d, num_labels=len(label_list), dropout_rate=out_dropout_rate)\n",
    "\n",
    "if torch.cuda.is_available():    \n",
    "    generator.to(device)\n",
    "    discriminator.to(device)\n",
    "    transformer.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "answering-sandwich",
   "metadata": {},
   "source": [
    "### Start Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fa50039",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "training_stats = []\n",
    "\n",
    "total_t0 = time.time()\n",
    "\n",
    "transformer_vars = [i for i in transformer.parameters()]\n",
    "d_vars = transformer_vars + [v for v in discriminator.parameters()]\n",
    "g_vars = [v for v in generator.parameters()]\n",
    "\n",
    "dis_optimizer = torch.optim.AdamW(d_vars, lr=learning_rate_discriminator)\n",
    "gen_optimizer = torch.optim.AdamW(g_vars, lr=learning_rate_generator) \n",
    "\n",
    "if apply_scheduler:\n",
    "    num_train_examples = len(train_examples)\n",
    "    num_train_steps = int(num_train_examples / batch_size * num_train_epochs)\n",
    "    num_warmup_steps = int(num_train_steps * warmup_proportion)\n",
    "\n",
    "    scheduler_d = get_constant_schedule_with_warmup(dis_optimizer, \n",
    "                                           num_warmup_steps = num_warmup_steps)\n",
    "    scheduler_g = get_constant_schedule_with_warmup(gen_optimizer, \n",
    "                                           num_warmup_steps = num_warmup_steps)\n",
    "    \n",
    "import torch\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "writer = SummaryWriter(log_dir=f\"{CHECKPOINT_DIR}/runs/\")\n",
    "\n",
    "step_cnt = 0 # to decide generator's/discriminator's backprop?\n",
    "best_f1 = 0\n",
    "best_checkpoint_so_far = None\n",
    "\n",
    "for epoch_i in range(0, num_train_epochs):\n",
    "    # ========================================\n",
    "    #               Training\n",
    "    # ========================================\n",
    "    # Perform one full pass over the training set.\n",
    "    print(\"\")\n",
    "    print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, num_train_epochs))\n",
    "    print('Training...')\n",
    "\n",
    "    # Measure how long the training epoch takes.\n",
    "    t0 = time.time()\n",
    "\n",
    "    # Reset the total loss for this epoch.\n",
    "    tr_g_loss = 0\n",
    "    tr_d_loss = 0\n",
    "\n",
    "    # Put the model into training mode.\n",
    "    transformer.train() \n",
    "\n",
    "    # For each batch of training data...\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "\n",
    "        step_cnt += 1\n",
    "        if step_cnt % 3 == 0:\n",
    "            generator.train()\n",
    "            discriminator.eval()\n",
    "        else:\n",
    "            generator.eval()\n",
    "            discriminator.train()\n",
    "        \n",
    "        # Progress update every print_each_n_step batches.\n",
    "        if step % print_each_n_step == 0 and not step == 0:\n",
    "            # Calculate elapsed time in minutes.\n",
    "            elapsed = format_time(time.time() - t0)\n",
    "            \n",
    "            # Report progress.\n",
    "            print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_dataloader), elapsed))\n",
    "\n",
    "        # Unpack this training batch from our dataloader. \n",
    "        b_input_ids = batch[0].to(device)\n",
    "        b_input_mask = batch[1].to(device)\n",
    "        b_labels = batch[2].to(device)\n",
    "        b_label_mask = batch[3].to(device)\n",
    "        b_labels01 = batch[4].to(device)\n",
    "        real_batch_size = b_input_ids.shape[0]\n",
    "     \n",
    "        model_outputs = transformer(b_input_ids, attention_mask=b_input_mask)\n",
    "        hidden_states = model_outputs[0]\n",
    "        \n",
    "\n",
    "        # Generator using Transformer Decoder \n",
    "        noise = torch.randn(max_seq_length, real_batch_size, hidden_size).to(device)\n",
    "        memory = torch.randn(max_seq_length, real_batch_size, hidden_size).to(device)\n",
    "        gen_rep = generator(noise, memory).permute(1, 0, 2)\n",
    "\n",
    "        # Generate the output of the Discriminator for real and fake data.\n",
    "        disciminator_input = torch.cat([hidden_states, gen_rep], dim=0)\n",
    "        \n",
    "        \n",
    "        # Then, we select the output of the disciminator\n",
    "        features, logits, probs = discriminator(disciminator_input)\n",
    "\n",
    "        # Finally, we separate the discriminator's output for the real and fake data\n",
    "        features_list = torch.split(features, real_batch_size)\n",
    "        D_real_features = features_list[0]\n",
    "        D_fake_features = features_list[1]\n",
    "      \n",
    "        logits_list = torch.split(logits, real_batch_size)\n",
    "        D_real_logits = logits_list[0]\n",
    "        D_fake_logits = logits_list[1]\n",
    "        \n",
    "        probs_list = torch.split(probs, real_batch_size)\n",
    "        D_real_probs = probs_list[0] # (5,64,3)\n",
    "        D_fake_probs = probs_list[1] # (5,64,3)\n",
    "        \n",
    "        \n",
    "        # Generator's LOSS estimation\n",
    "        g_loss_d = -1 * torch.mean(torch.log(1 - D_fake_probs[:, :, -1] + epsilon))\n",
    "        g_feat_reg = torch.mean(torch.pow(torch.mean(D_real_features, dim=0) - torch.mean(D_fake_features, dim=0), 2))\n",
    "        g_loss = g_loss_d + g_feat_reg\n",
    "\n",
    "        # Disciminator's LOSS estimation\n",
    "        logits = D_real_logits[:, :, 0:-1]\n",
    "        \n",
    "        \n",
    "        log_probs = F.log_softmax(logits, dim=-1)\n",
    "        label2one_hot = torch.nn.functional.one_hot(b_labels01, len(label_list))\n",
    "        per_example_loss = -torch.sum(label2one_hot * log_probs, dim=(-2, -1))\n",
    "        per_example_loss = torch.masked_select(per_example_loss, b_label_mask.to(device))\n",
    "        labeled_example_count = per_example_loss.type(torch.float32).numel()\n",
    "\n",
    "        # It may be the case that a batch does not contain labeled examples, so the \"supervised loss\" in this case is not evaluated\n",
    "        if labeled_example_count == 0:\n",
    "            D_L_Supervised = 0\n",
    "        else:\n",
    "            D_L_Supervised = torch.div(torch.sum(per_example_loss.to(device)), labeled_example_count)\n",
    "                 \n",
    "        D_L_unsupervised1U = -1 * torch.mean(torch.log(1 - D_real_probs[:, :, -1] + epsilon))\n",
    "        D_L_unsupervised2U = -1 * torch.mean(torch.log(D_fake_probs[:, :, -1] + epsilon))\n",
    "        d_loss = D_L_Supervised + D_L_unsupervised1U + D_L_unsupervised2U\n",
    "\n",
    "        #---------------------------------\n",
    "        #  OPTIMIZATION\n",
    "        #---------------------------------\n",
    "        # Avoid gradient accumulation\n",
    "        gen_optimizer.zero_grad()\n",
    "        dis_optimizer.zero_grad()\n",
    "        \n",
    "\n",
    "        if step_cnt % 3 == 0:\n",
    "            g_loss.backward()\n",
    "            gen_optimizer.step()\n",
    "            if apply_scheduler:\n",
    "                scheduler_g.step()\n",
    "        else:\n",
    "            d_loss.backward() \n",
    "            dis_optimizer.step()\n",
    "            if apply_scheduler:\n",
    "                scheduler_d.step()\n",
    "\n",
    "        \n",
    "        # Save the losses to print them later\n",
    "        tr_g_loss += g_loss.item()\n",
    "        tr_d_loss += d_loss.item()\n",
    "\n",
    "\n",
    "    # Calculate the average loss over all of the batches.\n",
    "    avg_train_loss_g = tr_g_loss / len(train_dataloader)\n",
    "    avg_train_loss_d = tr_d_loss / len(train_dataloader)             \n",
    "    \n",
    "    # Measure how long this epoch took.\n",
    "    training_time = format_time(time.time() - t0)\n",
    "\n",
    "    print(\"\")\n",
    "    print(\"  Average training loss generator: {0:.3f}\".format(avg_train_loss_g))\n",
    "    print(\"  Average training loss discriminator: {0:.3f}\".format(avg_train_loss_d))\n",
    "    print(\"  Training epoch took: {:}\".format(training_time))\n",
    "        \n",
    "        \n",
    "    # ========================================\n",
    "    #     TEST ON THE EVALUATION DATASET\n",
    "    # ========================================\n",
    "    print(\"\")\n",
    "    print(\"Running Validation...\")\n",
    "\n",
    "    t0 = time.time()\n",
    "\n",
    "    # Put the model in evaluation mode--the dropout layers behave differently during evaluation.\n",
    "    transformer.eval()\n",
    "    discriminator.eval()\n",
    "    generator.eval()\n",
    "\n",
    "    nll_loss = torch.nn.CrossEntropyLoss(ignore_index=-1)\n",
    "    total_valid_loss = 0\n",
    "    \n",
    "    all_preds = []\n",
    "    all_labels_ids = []\n",
    "\n",
    "    for batch in valid_dataloader:\n",
    "        \n",
    "        # Unpack this training batch from our dataloader. \n",
    "        b_input_ids = batch[0].to(device)\n",
    "        b_input_mask = batch[1].to(device)\n",
    "        b_labels = batch[2].to(device)\n",
    "        b_labels01 = batch[4].to(device)\n",
    "\n",
    "        with torch.no_grad():        \n",
    "            model_outputs = transformer(b_input_ids, attention_mask=b_input_mask)\n",
    "            hidden_states = model_outputs[0]\n",
    "            _, logits, probs = discriminator(hidden_states)\n",
    "            filtered_logits = logits[:, :, 0:-1]\n",
    "            \n",
    "            total_valid_loss += nll_loss(filtered_logits.reshape((-1, len(label_list))), b_labels01.reshape((-1)))\n",
    "        \n",
    "        _, preds = torch.max(filtered_logits, 2)\n",
    "        \n",
    "        all_preds += preds.detach().cpu()\n",
    "        all_labels_ids += b_labels.detach().cpu()\n",
    "\n",
    "    # Report the final accuracy for this validation run.\n",
    "    all_preds = torch.stack(all_preds).numpy()\n",
    "    all_labels_ids = torch.stack(all_labels_ids).numpy()\n",
    "    \n",
    "    pad_ids = np.full((all_preds.shape), PAD_VALUE)\n",
    "    valid_accuracy = np.sum(np.logical_and(all_preds == all_labels_ids, all_labels_ids != pad_ids)) / np.sum(all_labels_ids != pad_ids)\n",
    "    print(\"  Validation Accuracy: {0:.3f}\".format(valid_accuracy))\n",
    "    \n",
    "    # NOTE: We consider prediction = 2 (PAD token) as FLUENT\n",
    "    \n",
    "    tp = np.sum(np.logical_and(all_preds == all_labels_ids, all_labels_ids == np.full((all_preds.shape), 1)))\n",
    "    tn = np.sum(np.logical_and(all_labels_ids == np.full((all_preds.shape), 0), \n",
    "                               np.logical_or(all_preds == np.full((all_preds.shape), 0), \n",
    "                                             all_preds == np.full((all_preds.shape), 2))))\n",
    "    fn = np.sum(np.logical_and(all_labels_ids == np.full((all_preds.shape), 1), \n",
    "                               all_labels_ids != all_preds))\n",
    "    fp = np.sum(np.logical_and(all_preds == np.full((all_preds.shape), 1), \n",
    "                               all_labels_ids == np.full((all_preds.shape), 0)))\n",
    "    \n",
    "    precision = tp / (tp + fp)\n",
    "    recall = tp / (tp + fn)\n",
    "    f1 = 2 * precision * recall / (precision + recall)\n",
    "\n",
    "    # Calculate the average loss over all of the batches.\n",
    "    avg_valid_loss = total_valid_loss / len(valid_dataloader)\n",
    "    avg_valid_loss = avg_valid_loss.item()\n",
    "    \n",
    "    # Measure how long the validation run took.\n",
    "    valid_time = format_time(time.time() - t0)\n",
    "\n",
    "    print(\"  Validation Precision: {0:.3f}\".format(precision))\n",
    "    print(\"  Validation Recall: {0:.3f}\".format(recall))\n",
    "    print(\"  Validation F1 Score: {0:.3f}\".format(f1))\n",
    "    print(\"  Validation Loss: {0:.3f}\".format(avg_valid_loss))\n",
    "    print(\"  Validation took: {:}\".format(valid_time))\n",
    "\n",
    "    # Record all statistics from this epoch.\n",
    "    training_stats.append(\n",
    "        {\n",
    "            'epoch': epoch_i + 1,\n",
    "            'Training Loss generator': avg_train_loss_g,\n",
    "            'Training Loss discriminator': avg_train_loss_d,\n",
    "            'Training Time': training_time,\n",
    "            \n",
    "            'Validation Loss': avg_valid_loss,\n",
    "            'Validation Time': valid_time,\n",
    "            'Validation tp': tp.item(),\n",
    "            'Validation tn': tn.item(),\n",
    "            'Validation fp': fp.item(),\n",
    "            'Validation fn': fn.item(),\n",
    "            'Validation Accuracy': valid_accuracy.item(),\n",
    "            'Validation Precision': precision.item(),\n",
    "            'Validation Recall': recall.item(),\n",
    "            'Validation F1 Score': f1.item()\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    writer.add_scalar('Train Loss Generator', avg_train_loss_g, epoch_i)\n",
    "    writer.add_scalar('Train Loss Discriminator', avg_train_loss_d, epoch_i)\n",
    "    writer.add_scalar('Validation Loss', avg_valid_loss, epoch_i)\n",
    "    writer.add_scalar('Validation Accuracy', valid_accuracy, epoch_i)\n",
    "    writer.add_scalar('Validation Precision', precision, epoch_i)\n",
    "    writer.add_scalar('Validation Recall', recall, epoch_i)\n",
    "    writer.add_scalar('Validation F1 Score', f1, epoch_i)\n",
    "    writer.flush()\n",
    "    \n",
    "    if math.isnan(f1):\n",
    "        f1 = 0\n",
    "\n",
    "    if f1 >= best_f1 and f1 > 0:\n",
    "        os.system(f'mkdir -p {CHECKPOINT_DIR}/checkpoint-{epoch_i + 1}/')\n",
    "        torch.save(discriminator, f\"{CHECKPOINT_DIR}/checkpoint-{epoch_i + 1}/discriminator.pt\")\n",
    "        torch.save(generator, f\"{CHECKPOINT_DIR}/checkpoint-{epoch_i + 1}/generator.pt\")\n",
    "        transformer.save_pretrained(f\"{CHECKPOINT_DIR}/checkpoint-{epoch_i + 1}\")\n",
    "        tokenizer.save_pretrained(f\"{CHECKPOINT_DIR}/checkpoint-{epoch_i + 1}/\")\n",
    "\n",
    "        # Remove the prev BEST checkpoint\n",
    "        if f1 > best_f1 and best_checkpoint_so_far:\n",
    "            os.system(f\"rm -r {CHECKPOINT_DIR}/checkpoint-{best_checkpoint_so_far}/\")\n",
    "\n",
    "        best_f1 = f1\n",
    "        best_checkpoint_so_far = epoch_i + 1\n",
    "\n",
    "    print(\"Best checkpoint so far:\", best_checkpoint_so_far)\n",
    "\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f422a5fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(training_stats)):\n",
    "    if math.isnan(training_stats[i]['Validation F1 Score']):\n",
    "        # print(training_stats[i]['Validation F1 Score'])\n",
    "        training_stats[i]['Validation F1 Score'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "editorial-algorithm",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Best Validation performance\")\n",
    "print(\"Accuracy: {:0.2f}\".format(training_stats[0]['Validation Accuracy'] * 100))\n",
    "print(\"Precision: {:0.2f}\".format(training_stats[0]['Validation Precision'] * 100))\n",
    "print(\"Recall: {:0.2f}\".format(training_stats[0]['Validation Recall'] * 100))\n",
    "print(\"F1 Score: {:0.2f}\".format(training_stats[0]['Validation F1 Score'] * 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "close-sociology",
   "metadata": {},
   "source": [
    "### Load best model from training & infer on test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b6bafee",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_stats.sort(key = lambda x: (x['Validation Loss'], -x['Validation F1 Score'], -x['Validation Accuracy']))\n",
    "training_stats.sort(key = lambda x: (-x['Validation F1 Score'], -x['Validation Accuracy']))\n",
    "\n",
    "best_checkpoint = training_stats[0]['epoch']\n",
    "print(\"Best checkpoint is:\", best_checkpoint)\n",
    "\n",
    "model_dir = f\"{CHECKPOINT_DIR}/checkpoint-{best_checkpoint}\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "transformer = AutoModel.from_pretrained(model_dir).to(device)\n",
    "discriminator = torch.load(f\"{model_dir}/discriminator.pt\")\n",
    "\n",
    "# Uncomment next two cells if you want to load a model directly for inference on a defined test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hispanic-convertible",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_dir = \"/home/development/vineet/DDP_1/seq-gan-bert/checkpoints/presto-english-small-210523-7/checkpoint-27/\"\n",
    "# tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "# transformer = AutoModel.from_pretrained(model_dir).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "another-yahoo",
   "metadata": {},
   "outputs": [],
   "source": [
    "# discriminator = torch.load(f\"{model_dir}/discriminator.pt\")\n",
    "# generator = torch.load(f\"{model_dir}/generator.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e75dd4ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "def test_tokenizer(examples):\n",
    "    tokenized_inputs = tokenizer(examples[\"disfluent\"], truncation=True, max_length=512, is_split_into_words=True)\n",
    "    return tokenized_inputs\n",
    "\n",
    "\n",
    "# Evaluate on blind sentences\n",
    "test_dict = {\n",
    "                'disfluent': [sentence.split(\"\\t\")[0].split() for sentence in open(test_filename, 'r').readlines()],\n",
    "                'labels': [sentence.split(\"\\t\")[1].split() for sentence in open(test_filename, 'r').readlines()],\n",
    "            }\n",
    "\n",
    "test_dataset = Dataset.from_dict(test_dict)\n",
    "test_dataset = test_dataset.map(test_tokenizer, batched=True)\n",
    "tokenized_for_word_ids = tokenizer(test_dataset[\"disfluent\"], truncation=True, is_split_into_words=True)\n",
    "\n",
    "all_preds = []\n",
    "all_labels_ids = []\n",
    "\n",
    "transformer.eval()\n",
    "discriminator.eval()\n",
    "generator.eval()\n",
    "\n",
    "for batch in test_dataloader:\n",
    "        # Unpack this training batch from our dataloader.\n",
    "        b_input_ids = batch[0].to(device)\n",
    "        b_input_mask = batch[1].to(device)\n",
    "        b_labels = batch[2].to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            model_outputs = transformer(b_input_ids, attention_mask=b_input_mask)\n",
    "            hidden_states = model_outputs[0]\n",
    "            _, logits, probs = discriminator(hidden_states)\n",
    "            filtered_logits = logits[:, :, 0:-1]\n",
    "\n",
    "        _, preds = torch.max(filtered_logits, 2)\n",
    "        all_preds += preds.detach().cpu()\n",
    "        all_labels_ids += b_labels.detach().cpu()\n",
    "\n",
    "\n",
    "tp, fp, tn, fn = 0, 0, 0, 0\n",
    "\n",
    "for i in range(len(all_preds)):\n",
    "    actual_input = test_dataset[\"disfluent\"][i]\n",
    "    word_ids = tokenized_for_word_ids.word_ids(i)[1:-1] # Remove [CLS] & [SEP] token's word_id\n",
    "\n",
    "    print(\"Input           :\\t\", ' '.join(actual_input))\n",
    "    print(\"Ref Labels      :\\t\", ' '.join(map(str, all_labels_ids[i][1: 1 + len(word_ids)].tolist())))\n",
    "    print(\"Predicted Labels:\\t\", ' '.join(map(str, all_preds[i][1: 1 + len(word_ids)].tolist())))\n",
    "\n",
    "    for ref_label, pred_label in zip(all_labels_ids[i][1: 1 + len(word_ids)].tolist(), all_preds[i][1: 1 + len(word_ids)].tolist()):\n",
    "        if ref_label == pred_label:\n",
    "            if ref_label == 1:\n",
    "                tp += 1\n",
    "            else:\n",
    "                tn += 1\n",
    "        elif ref_label == 1:\n",
    "            fn += 1 # pred_label can be 0 or 2\n",
    "        elif ref_label == 0:\n",
    "            if pred_label == 1:\n",
    "                fp += 1\n",
    "            elif pred_label == 2:\n",
    "                tp += 1\n",
    "\n",
    "    # print(tokenized_input)\n",
    "    # print(word_ids)\n",
    "    # print(predictions[i][1:1 + len(tokenized_input)])\n",
    "\n",
    "    previous_word_idx = None\n",
    "    disfluent = 0 # Count of (predicted) disfluent subwords of a word\n",
    "    fluent = 0 # count of (predicted) fluent subwords of a words\n",
    "    fluent_sentence = []\n",
    "\n",
    "    for idx, prediction in enumerate(all_preds[i][1:1 + len(word_ids)]): # Remove [CLS] & [SEP] & PAD TOKEN predictions\n",
    "\n",
    "        # We add/ignore the previous word (based on how many subwords of the word were predicted disfluent).\n",
    "        # Added if count(fluent subwords) >= count(disfluent subwords)\n",
    "        if word_ids[idx] != previous_word_idx:\n",
    "            if previous_word_idx is not None and fluent >= disfluent:\n",
    "                fluent_sentence.append(actual_input[previous_word_idx])\n",
    "\n",
    "            fluent, disfluent = 0, 0\n",
    "\n",
    "        if prediction == 0 or prediction == 2: # consider prediction = 2 (PAD token) is also FLUENT\n",
    "            fluent += 1\n",
    "        else:\n",
    "            disfluent += 1\n",
    "\n",
    "        previous_word_idx = word_ids[idx]\n",
    "\n",
    "    # Don't forget to add the last word\n",
    "    if previous_word_idx is not None and fluent >= disfluent:\n",
    "        fluent_sentence.append(actual_input[previous_word_idx])\n",
    "\n",
    "    print(\"Prediction:\", ' '.join(fluent_sentence))\n",
    "\n",
    "    if i % 500 == 0 and i > 0:\n",
    "        print(f\"Testing Example-{i}\")\n",
    "\n",
    "\n",
    "print('tp, fp, fn, tn:', tp, fp, fn, tn)\n",
    "precision = 100 * tp / (tp + fp)\n",
    "recall = 100 * tp / (tp + fn)\n",
    "f1 = 2 * precision * recall / (precision + recall)\n",
    "\n",
    "print(\"Accuracy:\", round(100 * (tp + tn) / (tp + fp + fn + tn), 2))\n",
    "print(\"Precision:\", round(precision, 2))\n",
    "print(\"Recall:\", round(recall, 2))\n",
    "print(\"F1 Score:\", round(f1, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "funky-tomorrow",
   "metadata": {},
   "source": [
    "### Save data from training as log file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "secure-holocaust",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(training_stats[0])\n",
    "import json\n",
    "\n",
    "with open(f\"{CHECKPOINT_DIR}/log\", 'w') as log:\n",
    "    log.write(json.dumps(training_stats))"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}