train.py

# -*- coding: utf-8 -*-
"""My Experiments.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1owIZWIrCVhdZ_-mwPsfy8vW5DlyqXZE0

### **Hope Project**

#### Imports
---
"""

# !pip install -qq datasets pytorch-lightning comet-ml transformers optuna torchtext sacrebleu

from comet_ml import Experiment
from pytorch_lightning.loggers import CometLogger
import torch
from torch import Tensor, nn
from torch.utils.data import DataLoader, Dataset
import torch.nn.functional as F

import pytorch_lightning as pl
from pytorch_lightning import Trainer, Callback

from optuna.integration import PyTorchLightningPruningCallback
import optuna
from optuna.integration import PyTorchLightningPruningCallback
from optuna.visualization import plot_contour
from optuna.visualization import plot_intermediate_values
from optuna.visualization import plot_optimization_history
from optuna.visualization import plot_parallel_coordinate
from optuna.visualization import plot_param_importances
from optuna.visualization import plot_slice

# from datasets import  load_metric
from transformers import BertModel, BertTokenizerFast
from transformers.optimization import AdamW, get_linear_schedule_with_warmup
import numpy as np
import pandas as pd
import os
import regex
from nltk.tokenize import TweetTokenizer 
# from google.colab import drive
# drive.mount('/content/drive')

from hopfield import HopfieldLayer, Hopfield, HopfieldPooling

comet_logger = CometLogger(
    api_key="asUz6cXdIhbccxkk4Xr2ekODx",
    workspace="walexi",
    project_name="hope-detect",
    auto_param_logging=False,
    log_code=True,
    auto_metric_logging=False
)

"""#### Constants"""

EPOCHS = 3
TRIALS = 2
BATCH_SIZE = 8
PRE_TRAINED_MODEL_NAME = 'bert-base-cased'
CLASSES = {
    'english': ['Non_hope_speech', 'Hope_speech', 'not-English'],
    'tamil': ['Non_hope_speech', 'Hope_speech', 'not-Tamil'],
    'malayalam': ['Non_hope_speech', 'Hope_speech', 'not-Malayalam']
  }
MODEL_DIR = '/content'
DATA_ROOT ='/content/drive/MyDrive/Datasets' 
TRAIN_FILE = 'english_hope_train.csv'
VALIDATION_FILE = 'english_hope_dev.csv'
TEST_FILE = 'english_hope_test.csv'

"""#### Data Exploration

#### Datasets
"""

#  from https://github.com/ufoym/imbalanced-dataset-sampler/blob/master/torchsampler/imbalanced.py
class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler):
    """Samples elements randomly from a given list of indices for imbalanced dataset
    Arguments:
        indices (list, optional): a list of indices
        num_samples (int, optional): number of samples to draw
    """

    def __init__(self, dataset, indices=None, num_samples=None):
                
        # if indices is not provided, 
        # all elements in the dataset will be considered
        self.indices = list(range(len(dataset))) \
            if indices is None else indices

        # if num_samples is not provided, 
        # draw `len(indices)` samples in each iteration
        self.num_samples = len(self.indices) \
            if num_samples is None else num_samples
            
        # distribution of classes in the dataset 
        label_to_count = {}
        for idx in self.indices:
            label = self._get_label(dataset, idx)
            if label in label_to_count:
                label_to_count[label] += 1
            else:
                label_to_count[label] = 1
                
        # weight for each sample
        weights = [1.0 / label_to_count[self._get_label(dataset, idx)]
                   for idx in self.indices]
        self.weights = torch.DoubleTensor(weights)

    def _get_label(self, dataset, idx):
      return dataset[idx]['target'].argmax().item()
    def __iter__(self):
        return (self.indices[i] for i in torch.multinomial(
            self.weights, self.num_samples, replacement=True))

    def __len__(self):
        return self.num_samples

def read_csv(data_root, train_file, val_file, test_file=None):
  train_df = pd.read_table(data_root+'/'+train_file, index_col=False, names=['texts', 'labels'])
  val_df = pd.read_table(data_root+'/'+val_file, index_col=False, names=['texts', 'labels'])
  if(test_file):
    test_df = pd.read_table(data_root+'/'+test_file, index_col=False, names=['texts', 'labels'])
    return train_df, val_df, test_df
  else:
    return train_df, val_df

class HopeDataset(Dataset):
  def __init__(self, texts, labels, tokenizer, max_len, lang: str="english"):
    self.texts = texts
    self.labels = labels
    self.tweetTokenizer = None
    if isinstance(tokenizer, dict):
      self.tweetTokenizer = tokenizer['tweetTokenizer']
      self.tokenizer = tokenizer['tokenizer']
    else:
      self.tokenizer = tokenizer
    self.max_len = max_len
    self.lang = lang

  def one_hot(self, word: str):
    # a = np.zeros((1,3))
    # np.put(a, CLASSES.index(word), 1)

    return torch.zeros((1,3), dtype=torch.long).scatter(1, torch.tensor([[CLASSES[self.lang].index(word)]]), 1)
  
  def __len__(self):
    return len(self.texts)

  def __getitem__(self, item):
    text = self.texts[item]
    target = self.one_hot(self.labels[item])
    config = {
      'max_length' : self.max_len,
      'return_attention_mask' : True,
      'padding' : 'max_length',
      'truncation': True,
      'return_tensors' : 'pt',
      'return_token_type_ids' : False
    }

    if (self.tweetTokenizer):
      tokens = self.tweetTokenizer.tokenize(text)
      encoding = self.tokenizer(tokens, is_split_into_words=True, **config)
    else:
      encoding = self.tokenizer(text, **config)


    return {
        'text': text,
        'input_ids': encoding['input_ids'],
        'attention_mask' : encoding['attention_mask'],
        'target': target
    }

class HopeDataModule(pl.LightningDataModule):
  def __init__(self, batch_size: int=16, max_len: int=216, lang:str='english', data_root:str=DATA_ROOT, train_file: str=TRAIN_FILE, val_file: str=VALIDATION_FILE, test_file: str = TEST_FILE):
    super().__init__()
    self.bs = batch_size
    self.max_len = max_len
    self.lang = lang
    self.train_df, self.val_df = read_csv(data_root, train_file, val_file)
    self.train_len = len(self.train_df)

  
  def prepare_data(self):
    tokenizer = None
    tknz = BertTokenizerFast.from_pretrained(PRE_TRAINED_MODEL_NAME)
    if(self.lang=='english'):
      tokenizer = {}
      tweetTokenizer = TweetTokenizer(strip_handles=True, preserve_case=True, reduce_len=False)
      tokenizer['tweetTokenizer'] = tweetTokenizer
      tokenizer['tokenizer'] = tknz
    else:
      tokenizer = tknz
    
    self.train_set = HopeDataset(self.train_df.texts, self.train_df.labels, tokenizer, self.max_len)
    self.val_set = HopeDataset(self.val_df.texts, self.val_df.labels, tokenizer, self.max_len)
    
    
  
  # def setup(self, stage=None):
  #   if stage=='fit' or stage is None:
  #     self.train_set, self.val_set = self.dataset

  def train_dataloader(self):
    return DataLoader(self.train_set, batch_size=self.bs, sampler=ImbalancedDatasetSampler(self.train_set))
  
  def val_dataloader(self):
    return DataLoader(self.val_set, batch_size=self.bs)

  # def test_dataloader(self):
  #   return

"""##### Template"""

class BackBone(nn.Module):

  def __init__(self, num_classes, trials=None):

    super(BackBone, self).__init__()

    self.bert = BertModel.from_pretrained(PRE_TRAINED_MODEL_NAME)
    # self.bi_lstm = return sequences and apply dropout
    # max_pool
    # dense then dropout and dense_output
    self.hopfield_pool = HopfieldPooling(self.bert.config.hidden_size, hidden_size=32, output_size=self.bert.config.hidden_size, num_heads=1)
    if(trials):
      self.dropout = nn.Dropout(trials.suggest_uniform('dropout_rate', 0.2, 0.5))
    else:
      self.dropout = nn.Dropout(0.3)
    
    self.lin = nn.Linear(self.bert.config.hidden_size, num_classes)

  def forward(self, input_ids, attention_mask):
    output = self.bert(
        input_ids=input_ids, 
        attention_mask=attention_mask,
        output_attentions=True
        )
    out = output.pooler_output
    # print(out.size()) #8 768 
#     torch.Size([8, 768])
# torch.Size([1, 8, 768])
# torch.Size([1, 768])
    out = out.unsqueeze(1)
    # print(out.size())
    out = self.hopfield_pool(out)
    # print(out.size()) #1 768
    out = self.lin(self.dropout(out.squeeze(1)))
    return out, output.attentions

# num_heads = attention.shape[0]
# fig, axs = plt.subplot(3,4)
# for i in range(num_heads):

#   axs[i//4, i%4].imshow(attention[i], cmap='hot', interpolation='nearest')
# fig.show()

"""#### Lightning & Training"""

class Pipeline(pl.LightningModule):
  def __init__(self, model, total_steps:int, trial=None, max_len:int=100):
    super().__init__()

    self.trial = trial
    self.model = model
    self.max_len = max_len
    self.total_steps = total_steps
    # self.bleu_score = load_metric('sacrebleu') # make this generic and try a number of metric

  def cross_entropy_loss(self, logits, labels):
    loss =  nn.CrossEntropyLoss()
    return loss(logits, labels)
  # 'input_ids': encoding['input_ids'],
  #       'attention_mask' : encoding['attention_mask'],
  #       'target': target
  def training_step(self, batch, batch_idx):
    input_ids, attention_mask, target = batch['input_ids'], batch['attention_mask'], batch['target']

    logits, attention = self.model(input_ids.squeeze(1), attention_mask.squeeze(1))
    
    loss = self.cross_entropy_loss(logits, target.squeeze(1).argmax(dim=1))

    # self.bleu_score.add_batch(model_output, target)

    # logs = {'train_loss': loss.detach(), 'train_attention': attention[0].detach().cpu().numpy(), 'trial_num': -1}
    self.log('train_loss', loss.detach())
    # self.logger.experiment.log_metric('train_loss', loss)

    return loss

  # def training_epoch_end(self, training_step_outputs):
  #   #compute at the end of epoch, and log to comet
  #   # score = self.bleu_score.compute()
  #   pass

  def validation_step(self, batch, batch_idx):
    input_ids, attention_mask, target = batch['input_ids'], batch['attention_mask'], batch['target']

    # print(input_ids.size())
    logits, attention = self.model(input_ids.squeeze(1), attention_mask.squeeze(1))

    loss = self.cross_entropy_loss(logits, target.squeeze(1).argmax(dim=1))
    
    self.log('val_loss', loss.detach())
    # self.logger.experiment.log_metric('val_loss', loss)

    return loss

  # def validation_epoch_end(self, val_step_outputs):
  #   # randomly log the attention of a sample output
  #   # use the step with the lowest loss and log attention
  #   pass

  def configure_optimizers(self):
    if(self.trial):
      lr = self.trial.suggest_loguniform('learning_rate', 2e-5, 5e-5)
      self.logger.log_hyperparams({'lr': lr})
      optimizer = AdamW(self.parameters(), lr=lr, correct_bias=False)
    else:
      optimizer = AdamW(self.parameters(), lr=2e-5, correct_bias=False)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=self.total_steps)
    return [optimizer], [scheduler]

"""##### model params"""

# model.named_parameters

"""#### hyperparameter search"""

class MetricsCallback(Callback):
    """PyTorch Lightning metric callback."""

    def __init__(self):
      super().__init__()
      self.metrics = []

    def on_validation_end(self, trainer, pl_module):
      print(trainer.callback_metrics)
      self.metrics.append(trainer.callback_metrics)

def objective(trial, model_dir: str=MODEL_DIR, epochs: int=EPOCHS, batch_size:int=BATCH_SIZE, data_root: str=DATA_ROOT, train_file: str=TRAIN_FILE, val_file: str=VALIDATION_FILE, test_file:str=TEST_FILE):
  # Filenames for each trial must be made unique in order to access each checkpoint.
  checkpoint_callback = pl.callbacks.ModelCheckpoint(
        os.path.join(model_dir, "trial_{}".format(trial.number), "{epoch}"), monitor="val_loss", mode='min'
    )
  metrics_callback = MetricsCallback()

  trainer = Trainer(
      logger=comet_logger,
      gpus=-1 if torch.cuda.is_available() else None,
      checkpoint_callback=checkpoint_callback,
      # precision=16, 
      progress_bar_refresh_rate=0, 
      max_epochs=EPOCHS,
      callbacks = [metrics_callback, PyTorchLightningPruningCallback(trial, monitor="val_loss")]
      )

  dataloader = HopeDataModule(batch_size=batch_size, data_root=data_root, train_file=train_file, val_file=val_file, test_file = test_file)

  model = BackBone(len(CLASSES['english']), trials=trial)
  total_steps = dataloader.train_len * EPOCHS
  pipeline = Pipeline(model, total_steps, trial)
  trainer.fit(pipeline, datamodule=dataloader)

  return metrics_callback.metrics[-1]["val_loss"].item()


if __name__ == '__main__':
   # check_git_status()

    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=4)
    parser.add_argument('--batch-size', type=int, default=16)
    parser.add_argument('--model-dir', type=str, default=".")
    parser.add_argument('--dataset-root', type=str, default=".")
    parser.add_argument('--train-file', type=str, default='data/debris.yaml', help='*.data path')
    parser.add_argument('--val-file', type=str, default='data/debris.yaml', help='*.data path')
    parser.add_argument('--test-file', type=str, default='data/debris.yaml', help='*.data path')
    parser.add_argument('--pruning', action='store_true', help='cache images for faster training')


    args = parser.parse_args()

    pruner = optuna.pruners.MedianPruner() if args.pruning else optuna.pruners.NopPruner()

    study = optuna.create_study(direction="maximize", pruner=pruner)
    study.optimize(objective(model_dir=args.model-dir,data_root=args.dataset-rooot, train_file=args.train-file,val_file=args.val-file), n_trials=100, timeout=600)

    print("Number of finished trials: {}".format(len(study.trials)))

    print("Best trial:")
    trial = study.best_trial

    print("  Value: {}".format(trial.value))

    print("  Params: ")
    for key, value in trial.params.items():
        print("    {}: {}".format(key, value))

    shutil.rmtree(MODEL_DIR)