train_generative.py

# Importing libraries
import datetime
import os
import sys
import time
import torch
from shutil import copyfile

import numpy as np
import pandas as pd
import spacy
import transformers
from rich import box
from rich.console import Console
from rich.table import Column, Table
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import MT5Tokenizer, MT5ForConditionalGeneration

import e2e_tbsa_preprocess
import evaluate_e2e_tbsa
import utils
from utils import EarlyStopping, YourDataSetClass

# pd.set_option('display.max_colwidth', -1)


# define a rich console logger
console = Console(record=True)

# Set random seeds and deterministic pytorch for reproducibility
torch.manual_seed(0)  # pytorch random seed
np.random.seed(0)  # numpy random seed
torch.backends.cudnn.deterministic = True

# Setting up the device for GPU usage
from torch import cuda

device = 'cuda' if cuda.is_available() else 'cpu'

LANG_MAP = {'en': 'english', 'es': 'spanish', 'ru': 'russian'}
SEED_LIST = [0, 1, 2]
LR_LIST = [5e-5]

if sys.argv[1] == "false":
    FULL_DATASET = False
else:
    FULL_DATASET = True


if len(sys.argv) == 3:
    assert sys.argv[2][0] == '[' and sys.argv[2][-1] == ']'
    SEED_LIST = sys.argv[2][1:-1]
    SEED_LIST = SEED_LIST.split(sep='_')
    SEED_LIST = [int(x) for x in SEED_LIST]

print("SEEDS: {}".format(SEED_LIST))


def train(tokenizer, model, device, loader, optimizer):
    """
    Function to be called for training with the parameters passed from main function
    """
    train_losses = []
    model.train()
    for _, data in tqdm(enumerate(loader, 0), total=len(loader), desc='Processing batches..'):
        y = data['target_ids'].to(device, dtype=torch.long)
        lm_labels = y.clone()
        lm_labels[y == tokenizer.pad_token_id] = -100
        ids = data['source_ids'].to(device, dtype=torch.long)
        mask = data['source_mask'].to(device, dtype=torch.long)

        outputs = model(input_ids=ids, attention_mask=mask, labels=lm_labels)
        loss = outputs[0]

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_losses.append(loss.item())
    return train_losses


def validate(tokenizer, model, device, loader):
    """
    Function to be called for validating the trainner with the parameters passed from main function
    """
    validate_losses = []
    model.eval()
    for _, data in tqdm(enumerate(loader, 0), total=len(loader), desc='Validating batches..'):
        y = data['target_ids'].to(device, dtype=torch.long)
        lm_labels = y.clone()
        lm_labels[y == tokenizer.pad_token_id] = -100
        ids = data['source_ids'].to(device, dtype=torch.long)
        mask = data['source_mask'].to(device, dtype=torch.long)
        outputs = model(input_ids=ids, attention_mask=mask, labels=lm_labels)
        loss = outputs[0]
        validate_losses.append(loss.item())
    return validate_losses


def build_data(model_params, dataframes, source_text, target_text):
    # tokenzier for encoding the text
    tokenizer = MT5Tokenizer.from_pretrained(model_params["MODEL"])
    tokenizer.add_tokens(['<sep>', '<lang>'])  # , 'generate_english', 'generate_spanish', 'generate_russian'])

    # logging
    console.log(f"[Data]: Reading data...\n")

    # Creation of Dataset and Dataloader
    train_dataset = dataframes[0].sample(frac=1, random_state=model_params["SEED"]).reset_index(drop=True)
    val_dataset = dataframes[1].reset_index(drop=True)
    test_dataset = dataframes[2].reset_index(drop=True)
    console.print(f"TRAIN Dataset: {train_dataset.shape}")
    console.print(f"VALIDATION Dataset: {val_dataset.shape}")
    console.print(f"TEST Dataset: {test_dataset.shape}\n")

    # Creating the Training and Validation dataset for further creation of Dataloader
    training_set = YourDataSetClass(train_dataset, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"],
                                    model_params["MAX_TARGET_TEXT_LENGTH"], source_text, target_text)
    val_set = YourDataSetClass(val_dataset, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"],
                               model_params["MAX_TARGET_TEXT_LENGTH"], source_text, target_text)
    test_set = YourDataSetClass(test_dataset, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"],
                                model_params["MAX_TARGET_TEXT_LENGTH"], source_text, target_text)

    # Defining the parameters for creation of dataloaders
    train_params = {'batch_size': model_params["TRAIN_BATCH_SIZE"], 'shuffle': True, 'num_workers': 2}
    val_params = {'batch_size': model_params["VALID_BATCH_SIZE"], 'shuffle': False, 'num_workers': 2}
    test_params = {'batch_size': model_params["VALID_BATCH_SIZE"], 'shuffle': False, 'num_workers': 2}

    # Creation of Dataloaders for testing and validation. This will be used down for training and validation stage for the model.
    training_loader = DataLoader(training_set, **train_params)
    validation_loader = DataLoader(val_set, **val_params)
    test_loader = DataLoader(test_set, **test_params)

    return training_loader, validation_loader, test_loader, tokenizer


def generate(tokenizer, model, device, loader, model_params):
    """
  Function to evaluate model for spanbert-predictions

  """
    model.eval()
    predictions = []
    actuals = []
    data_list = []
    with torch.no_grad():
        for _, data in enumerate(loader, 0):
            y = data['target_ids'].to(device, dtype=torch.long)
            ids = data['source_ids'].to(device, dtype=torch.long)
            mask = data['source_mask'].to(device, dtype=torch.long)

            generated_ids = model.generate(,

            preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in
                     generated_ids]
            target = [tokenizer.decode(t, skip_special_tokens=True, clean_up_tokenization_spaces=True) for t in y]
            if _ % 10 == 0:
                console.print(f'Completed {_}')

            predictions.extend(preds)
            actuals.extend(target)
            data_list.extend(data["sentences_texts"])
    return predictions, actuals, data_list


def T5Trainer(training_loader, validation_loader, tokenizer, model_params):
    """
    T5 trainer
    """

    torch.manual_seed(model_params['SEED'])  # pytorch random seed
    np.random.seed(model_params['SEED'])  # numpy random seed

    # logging
    console.log(f"""[Model]: Loading {model_params["MODEL"]}...\n""")

    # training_logger = Table(Column("Epoch", justify="center"), Column("train_loss", justify="center"),
    #                         Column("val_loss", justify="center"), Column("Epoch Time", justify="center"),
    #                         title="Training Status", pad_edge=False, box=box.ASCII)

    training_logger = Table(Column("Epoch", justify="center"), Column("train_loss", justify="center"),
                            Column("val_f1", justify="center"), Column("Epoch Time", justify="center"),
                            title="Training Status", pad_edge=False, box=box.ASCII)

    # Defining the model. We are using t5-base model and added a Language model layer on top for generation of Summary.
    # Further this model is sent to device (GPU/TPU) for using the hardware.
    model = MT5ForConditionalGeneration.from_pretrained(model_params["MODEL"])
    model = model.to(device)

    # Defining the optimizer that will be used to tune the weights of the network in the training session.
    optimizer = transformers.Adafactor(params=model.parameters(), lr=model_params["LEARNING_RATE"],
                                       scale_parameter=False, relative_step=False)
    # initialize the early_stopping object
    early_stopping = EarlyStopping(patience=model_params["early_stopping_patience"], verbose=False, delta=0.1,
                                   path=f'{model_params["OUTPUT_PATH"]}/best_pytorch_model.bin')

    # Training loop
    console.log(f'[Initiating Fine Tuning]...\n')
    avg_train_losses = []
    # avg_valid_losses = []

    for epoch in range(model_params["TRAIN_EPOCHS"]):

        start_time = time.time()
        train_losses = train(tokenizer, model, device, training_loader, optimizer)
        # valid_losses = validate(tokenizer, model, device, validation_loader)
        epoch_time = round(time.time() - start_time)

        # calculate average loss over an epoch
        train_loss = np.average(train_losses)
        # valid_loss = np.average(valid_losses)
        avg_train_losses.append(train_loss)
        # avg_valid_losses.append(valid_loss)

        print("Early stopping: Calculating VALIDATION SCORE: ")
        prediction_file_name_validation = 'evaluation_predictions_val.csv'
        predictions_filepath_validation = '{}/{}'.format(model_params["OUTPUT_PATH"], prediction_file_name_validation)
        transformed_target_path_validation = '{}/{}'.format(model_params["OUTPUT_PATH"], "transformed_target_val.csv")
        transformed_sentiment_path_validation = '{}/{}'.format(model_params["OUTPUT_PATH"],
                                                               "transformed_sentiment_val.csv")

        T5Generator(validation_loader, model_params=model_params, output_file=prediction_file_name_validation,
                    model=model, tokenizer=tokenizer)

        valid_f1 = preprocess_and_evaluate(predictions_filepath_validation, "", transformed_sentiment_path_validation,
                                           transformed_target_path_validation, True)

        # preparing the processing time for the epoch and est. the total.
        epoch_time_ = str(datetime.timedelta(seconds=epoch_time))

        # total_time_estimated_ = str(
        #     datetime.timedelta(seconds=(epoch_time * (model_params["TRAIN_EPOCHS"] - epoch - 1))))
        # training_logger.add_row(f'{epoch + 1}/{model_params["TRAIN_EPOCHS"]}', f'{train_loss:.5f}', f'{valid_loss:.5f}',
        #                         f'{epoch_time_} (Total est. {total_time_estimated_})')
        total_time_estimated_ = str(
            datetime.timedelta(seconds=(epoch_time * (model_params["TRAIN_EPOCHS"] - epoch - 1))))
        training_logger.add_row(f'{epoch + 1}/{model_params["TRAIN_EPOCHS"]}', f'{train_loss:.5f}', f'{valid_f1}',
                                f'{epoch_time_} (Total est. {total_time_estimated_})')
        console.print(training_logger)

        # early_stopping needs the validation loss to check if it has decresed, 
        # and if it has, it will make a checkpoint of the current model
        # early_stopping(valid_loss, model)
        early_stopping(valid_f1, model)
        if early_stopping.early_stop:
            print("Early stopping")
            break

    console.log(f"[Saving Model]...\n")
    # Saving the model after training
    path = os.path.join(model_params["OUTPUT_PATH"], "model_files")
    model.save_pretrained(path)
    tokenizer.save_pretrained(path)
    console.log(f"[Replace best model with the last model]...\n")
    os.rename(f'{model_params["OUTPUT_PATH"]}/model_files/pytorch_model.bin',
              f'{model_params["OUTPUT_PATH"]}/model_files/last_epoch_pytorch_model.bin')
    copyfile(f'{model_params["OUTPUT_PATH"]}/best_pytorch_model.bin',
             f'{model_params["OUTPUT_PATH"]}/model_files/pytorch_model.bin')
    os.remove(f'{model_params["OUTPUT_PATH"]}/best_pytorch_model.bin')


def T5Generator(validation_loader, model_params, output_file, model=None, tokenizer=None):
    torch.manual_seed(model_params['SEED'])  # pytorch random seed
    np.random.seed(model_params['SEED'])  # numpy random seed

    console.log(f"[Loading Model]...\n")
    # Saving the model after training
    path = os.path.join(model_params["OUTPUT_PATH"], "model_files")

    if model is None:
        console.log("Using model from path")
        model = MT5ForConditionalGeneration.from_pretrained(path)
    else:
        console.log("Using existing model")

    if tokenizer is None:
        console.log("Using tokenizer from path")
        tokenizer = MT5Tokenizer.from_pretrained(path)
    else:
        console.log("Using existing tokenizer")

    model = model.to(device)

    # evaluating test dataset
    console.log(f"[Initiating Validation]...\n")

    predictions, actuals, data_list = generate(tokenizer, model, device, validation_loader, model_params)
    final_df = pd.DataFrame({'Generated Text': predictions, 'Actual Text': actuals, 'Original Sentence': data_list})
    final_df.to_csv(os.path.join(model_params["OUTPUT_PATH"], output_file))

    console.save_text(os.path.join(model_params["OUTPUT_PATH"], 'logs.txt'))

    console.log(f"[Validation Completed.]\n")
    console.print(f"""[Model] Model saved @ {os.path.join(model_params["OUTPUT_PATH"], "model_files")}\n""")
    console.print(
        f"""[Validation] Generation on Validation data saved @ {os.path.join(model_params["OUTPUT_PATH"], output_file)}\n""")
    console.print(f"""[Logs] Logs saved @ {os.path.join(model_params["OUTPUT_PATH"], 'logs.txt')}\n""")


def run_program_for_seed(seed, lr):
    MODEL_DIRECTORY = f"./generative-predictions_t5_full_{FULL_DATASET}_seed_{seed}/{train_domain}_{train_language}"

    model_params = {
        "OUTPUT_PATH": MODEL_DIRECTORY,  # output path
        "MODEL": "google/mt5-base",  # model_type: t5-base/t5-large
        "TRAIN_BATCH_SIZE": 8,  # training batch size
        "VALID_BATCH_SIZE": 8,  # validation batch size
        "TRAIN_EPOCHS": 50,  # number of training epochs
        "VAL_EPOCHS": 1,  # number of validation epochs
        "LEARNING_RATE": lr,  # learning rate
        "MAX_SOURCE_TEXT_LENGTH": 256,  # max length of source text
        "MAX_TARGET_TEXT_LENGTH": 64,  # max length of target text
        "early_stopping_patience": 10,  # number of epochs before stopping training.
        "SEED": seed
    }

    print(model_params)

    if not os.path.exists(MODEL_DIRECTORY):
        os.makedirs(MODEL_DIRECTORY)

    predictions_file = f'{test_domain}_{test_language}_predictions.csv'
    prediction_file_path = os.path.join(model_params["OUTPUT_PATH"], predictions_file)
    transformed_targets_file = f'{test_domain}_{test_language}_transformed-targets.csv'
    transformed_targets_file_path = os.path.join(model_params["OUTPUT_PATH"], transformed_targets_file)
    transformed_sentiments_file = f'{test_domain}_{test_language}_transformed-sentiments.csv'
    transformed_sentiments_file_path = os.path.join(model_params["OUTPUT_PATH"], transformed_sentiments_file)

    training_loader, validation_loader, test_loader, tokenizer = build_data(model_params,
                                                                            dataframes=[training, validation, test],
                                                                            source_text="sentences_texts",
                                                                            target_text="sentences_opinions")

    T5Trainer(training_loader, validation_loader, tokenizer, model_params=model_params)
    T5Generator(test_loader, model_params=model_params, output_file=predictions_file)

    preprocess_and_evaluate(prediction_file_path, seed, transformed_sentiments_file_path,
                            transformed_targets_file_path, False, lr)


def preprocess_and_evaluate(prediction_file_path, seed, transformed_sentiments_file_path,
                            transformed_targets_file_path, validation, lr=0):
    nlp = spacy.load(utils.get_spacy_language(test_language), disable=['parser', 'ner'])
    print(f"Preprocessing file at {prediction_file_path}")
    e2e_tbsa_preprocess.transform_gold_and_truth(test_language, nlp, prediction_file_path,
                                                 transformed_targets_file_path, transformed_sentiments_file_path)
    print(f"Evaluating files at {transformed_targets_file_path} and {transformed_sentiments_file_path}")
    predicted_data, gold_data = evaluate_e2e_tbsa.read_transformed_sentiments(transformed_sentiments_file_path)

    if validation is False:
        print(f"Evaluated test file at: {transformed_sentiments_file_path}")
        print(f"SEED: {seed}, LR: {lr}, TEST SET OUTPUT: {evaluate_e2e_tbsa.evaluate_ts(gold_data, predicted_data)}")
        print("\n--------------------------\n")
    else:
        return evaluate_e2e_tbsa.evaluate_ts(gold_data, predicted_data)[2]


if __name__ == '__main__':
    # domain: Rest16, Lap14, Mams, Mams_short
    # lang: en, es, ru
    for train_settings in [('Rest16', 'en', 'Rest16', 'en'), ('Rest16', 'es', 'Rest16', 'es'),
                           ('Lap14', 'en', 'Lap14', 'en'), ('Mams', 'en', 'Mams', 'en'),
                           ('Rest16', 'ru', 'Rest16', 'ru')]:

        train_domain = train_settings[0]
        train_language = train_settings[1]
        test_domain = train_settings[2]
        test_language = train_settings[3]

        if FULL_DATASET:
            training_file = './data/processed_full_train_{}_{}.csv'.format(train_domain, train_language)
        else:
            training_file = './data/processed_train_{}_{}.csv'.format(train_domain, train_language)

        validation_file = './data/processed_val_{}_{}.csv'.format(train_domain, train_language)
        test_file = './data/processed_test_{}_{}.csv'.format(test_domain, test_language)
        print("----------------\n\n"
              "Experiment: Training on {}.{}, Testing on {}.{}".format(train_domain, train_language, test_domain,
                                                                       test_language))

        training = pd.read_csv(training_file)
        validation = pd.read_csv(validation_file)
        test = pd.read_csv(test_file)

        # for cross-lingual
        # training['sentences_texts'] = training['sentences_texts'].map(lambda txt: f'generate {lang_map[train_settings[1]]} </s> {txt}')
        # validation['sentences_texts'] = validation['sentences_texts'].map(lambda txt: f'generate {lang_map[train_settings[1]]} </s> {txt}')
        # test['sentences_texts'] = test['sentences_texts'].map(lambda txt: f'generate {lang_map[train_settings[1]]} </s> {txt}')

        for lr in LR_LIST:
            for seed in SEED_LIST:
                run_program_for_seed(seed, lr)

    # test_language = 'en'
    # train_domain = ''
    # train_language = ''
    # test_domain = ''
    # training=None
    # validation=None
    # test = None
    #
    # preprocess_and_evaluate('/Users/dhruvmullick/Projects/GenerativeAspectBasedSentimentAnalysis/evaluation_predictions_val.csv', 0,
    #                         '/Users/dhruvmullick/Projects/GenerativeAspectBasedSentimentAnalysis/evaluation_predictions_val_sent.csv',
    #                         '/Users/dhruvmullick/Projects/GenerativeAspectBasedSentimentAnalysis/evaluation_predictions_val_asp.csv',False)