main.py

import argparse

from torch.utils.data import DataLoader
import torch
import torch.nn as nn

from src.bert import BERT
from src.pretrainer import BERTTrainer, BERTFineTuneTrainer, BERTAttention
from src.dataset import PretrainerDataset, TokenizerDataset
from src.vocab import Vocab

import time
import os
import tqdm
import pickle

def train():
    parser = argparse.ArgumentParser()

    parser.add_argument('-workspace_name', type=str, default=None)
    parser.add_argument('-code', type=str, default=None, help="folder for pretraining outputs and logs")
    parser.add_argument('-finetune_task', type=str, default=None, help="folder inside finetuning")
    parser.add_argument("-attention", type=bool, default=False, help="analyse attention scores")
    parser.add_argument("-diff_test_folder", type=bool, default=False, help="use for different test folder")
    parser.add_argument("-embeddings", type=bool, default=False, help="get and analyse embeddings")
    parser.add_argument('-embeddings_file_name', type=str, default=None, help="file name of embeddings")
    parser.add_argument("-pretrain", type=bool, default=False, help="pretraining: true, or false")
    # parser.add_argument('-opts', nargs='+', type=str, default=None, help='List of optional steps')
    parser.add_argument("-max_mask", type=int, default=0.15, help="% of input tokens selected for masking") 
    # parser.add_argument("-p", "--pretrain_dataset", type=str, default="pretraining/pretrain.txt", help="pretraining dataset for bert")
    # parser.add_argument("-pv", "--pretrain_val_dataset", type=str, default="pretraining/test.txt", help="pretraining validation dataset for bert")
# default="finetuning/test.txt",
    parser.add_argument("-vocab_path", type=str, default="pretraining/vocab.txt", help="built vocab model path with bert-vocab")

    parser.add_argument("-train_dataset_path", type=str, default="train.txt", help="fine tune train dataset for progress classifier")
    parser.add_argument("-val_dataset_path", type=str, default="val.txt", help="test set for evaluate fine tune train set")
    parser.add_argument("-test_dataset_path", type=str, default="test.txt", help="test set for evaluate fine tune train set")
    parser.add_argument("-num_labels", type=int, default=2, help="Number of labels") 
    parser.add_argument("-train_label_path", type=str, default="train_label.txt", help="fine tune train dataset for progress classifier")
    parser.add_argument("-val_label_path", type=str, default="val_label.txt", help="test set for evaluate fine tune train set")
    parser.add_argument("-test_label_path", type=str, default="test_label.txt", help="test set for evaluate fine tune train set")
    ##### change Checkpoint for finetuning
    parser.add_argument("-pretrained_bert_checkpoint", type=str, default=None, help="checkpoint of saved pretrained bert model")  #."output_feb09/bert_trained.model.ep40"
    parser.add_argument('-check_epoch', type=int, default=None)

    parser.add_argument("-hs", "--hidden", type=int, default=64, help="hidden size of transformer model") #64
    parser.add_argument("-l", "--layers", type=int, default=4, help="number of layers") #4
    parser.add_argument("-a", "--attn_heads", type=int, default=4, help="number of attention heads") #8
    parser.add_argument("-s", "--seq_len", type=int, default=50, help="maximum sequence length")

    parser.add_argument("-b", "--batch_size", type=int, default=500, help="number of batch_size") #64
    parser.add_argument("-e", "--epochs", type=int, default=50)#1501, help="number of epochs") #501
    # Use 50 for pretrain, and 10 for fine tune
    parser.add_argument("-w", "--num_workers", type=int, default=4, help="dataloader worker size")

    # Later run with cuda
    parser.add_argument("--with_cuda", type=bool, default=True, help="training with CUDA: true, or false")
    parser.add_argument("--log_freq", type=int, default=10, help="printing loss every n iter: setting n")
    # parser.add_argument("--corpus_lines", type=int, default=None, help="total number of lines in corpus")
    parser.add_argument("--cuda_devices", type=int, nargs='+', default=None, help="CUDA device ids")
    # parser.add_argument("--on_memory", type=bool, default=False, help="Loading on memory: true or false")

    parser.add_argument("--dropout", type=float, default=0.1, help="dropout of network")
    parser.add_argument("--lr", type=float, default=1e-05, help="learning rate of adam") #1e-3
    parser.add_argument("--adam_weight_decay", type=float, default=0.01, help="weight_decay of adam")
    parser.add_argument("--adam_beta1", type=float, default=0.9, help="adam first beta value")
    parser.add_argument("--adam_beta2", type=float, default=0.98, help="adam first beta value") #0.999

    parser.add_argument("-o", "--output_path", type=str, default="bert_trained.seq_encoder.model", help="ex)output/bert.model")    
    # parser.add_argument("-o", "--output_path", type=str, default="output/bert_fine_tuned.model", help="ex)output/bert.model")
    
    args = parser.parse_args()
    for k,v in vars(args).items():
        if 'path' in k:
            if v:
                if k == "output_path":
                    if args.code:
                        setattr(args, f"{k}", args.workspace_name+f"/output/{args.code}/"+v)
                    elif args.finetune_task:
                        setattr(args, f"{k}", args.workspace_name+f"/output/{args.finetune_task}/"+v)
                    else:
                        setattr(args, f"{k}", args.workspace_name+"/output/"+v)
                elif k != "vocab_path":
                    if args.pretrain:
                        setattr(args, f"{k}", args.workspace_name+"/pretraining/"+v)
                    else:
                        if args.code:
                            setattr(args, f"{k}", args.workspace_name+f"/{args.code}/"+v)
                        elif args.finetune_task:
                            if args.diff_test_folder and "test" in k:
                                setattr(args, f"{k}", args.workspace_name+f"/finetuning/"+v)
                            else:
                                setattr(args, f"{k}", args.workspace_name+f"/finetuning/{args.finetune_task}/"+v)
                        else:
                            setattr(args, f"{k}", args.workspace_name+"/finetuning/"+v)
                else:
                    setattr(args, f"{k}", args.workspace_name+"/"+v)
                
                print(f"args.{k} : {getattr(args, f'{k}')}")

    print("Loading Vocab", args.vocab_path)
    vocab_obj = Vocab(args.vocab_path)
    vocab_obj.load_vocab()
    print("Vocab Size: ", len(vocab_obj.vocab))
    
    if args.attention:
        print(f"Attention aggregate...... code: {args.code}, dataset: {args.finetune_task}")
        if args.code:
            new_folder = f"{args.workspace_name}/plots/{args.code}/"
            if not os.path.exists(new_folder):
                os.makedirs(new_folder)
        
        train_dataset = TokenizerDataset(args.train_dataset_path, None, vocab_obj, seq_len=args.seq_len)
        train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)
        print("Load Pre-trained BERT model")
        cuda_condition = torch.cuda.is_available() and args.with_cuda
        device = torch.device("cuda:0" if cuda_condition else "cpu")
        bert = torch.load(args.pretrained_bert_checkpoint, map_location=device)
        trainer = BERTAttention(bert, vocab_obj, train_dataloader = train_data_loader, workspace_name = args.workspace_name, code=args.code, finetune_task = args.finetune_task)
        trainer.getAttention()
        
    elif args.embeddings:
        print("Get embeddings... and cluster... ")
        train_dataset = TokenizerDataset(args.test_dataset_path, None, vocab_obj, seq_len=args.seq_len)
        train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)
        print("Load Pre-trained BERT model")
        cuda_condition = torch.cuda.is_available() and args.with_cuda
        device = torch.device("cuda:0" if cuda_condition else "cpu")
        bert = torch.load(args.pretrained_bert_checkpoint).to(device)
        available_gpus = list(range(torch.cuda.device_count()))
        if torch.cuda.device_count() > 1:
            print("Using %d GPUS for BERT" % torch.cuda.device_count())
            bert = nn.DataParallel(bert, device_ids=available_gpus)

        data_iter = tqdm.tqdm(enumerate(train_data_loader), 
                              desc="Model: %s" % (args.pretrained_bert_checkpoint.split("/")[-1]), 
                              total=len(train_data_loader), bar_format="{l_bar}{r_bar}")
        all_embeddings = []
        for i, data in data_iter:
            data = {key: value.to(device) for key, value in data.items()}
            embedding = bert(data["input"], data["segment_label"])
            # print(embedding.shape, embedding[:, 0].shape)
            embeddings = [h for h in embedding[:,0].cpu().detach().numpy()]
            all_embeddings.extend(embeddings)
            
        new_emb_folder = f"{args.workspace_name}/embeddings"
        if not os.path.exists(new_emb_folder):
            os.makedirs(new_emb_folder)
        pickle.dump(all_embeddings, open(f"{new_emb_folder}/{args.embeddings_file_name}.pkl", "wb"))
    else:
        if args.pretrain:
            print("Pre-training......")
            print("Loading Pretraining Train Dataset", args.train_dataset_path)
            print(f"Workspace: {args.workspace_name}")
            pretrain_dataset = PretrainerDataset(args.train_dataset_path, vocab_obj, seq_len=args.seq_len, max_mask = args.max_mask)

            print("Loading Pretraining Validation Dataset", args.val_dataset_path)
            pretrain_valid_dataset = PretrainerDataset(args.val_dataset_path, vocab_obj, seq_len=args.seq_len, max_mask = args.max_mask) \
                if args.val_dataset_path is not None else None
            
            print("Loading Pretraining Test Dataset", args.test_dataset_path)
            pretrain_test_dataset = PretrainerDataset(args.test_dataset_path, vocab_obj, seq_len=args.seq_len, max_mask = args.max_mask) \
                if args.test_dataset_path is not None else None

            print("Creating Dataloader")
            pretrain_data_loader = DataLoader(pretrain_dataset, batch_size=args.batch_size, num_workers=args.num_workers)
            pretrain_val_data_loader = DataLoader(pretrain_valid_dataset, batch_size=args.batch_size, num_workers=args.num_workers)\
                if pretrain_valid_dataset is not None else None
            pretrain_test_data_loader = DataLoader(pretrain_test_dataset, batch_size=args.batch_size, num_workers=args.num_workers)\
                if pretrain_test_dataset is not None else None
            
            print("Building BERT model")
            bert = BERT(len(vocab_obj.vocab), hidden=args.hidden, n_layers=args.layers, attn_heads=args.attn_heads, dropout=args.dropout)

            if args.pretrained_bert_checkpoint:
                print(f"BERT model : {args.pretrained_bert_checkpoint}")
                bert = torch.load(args.pretrained_bert_checkpoint)
            
            new_log_folder = f"{args.workspace_name}/logs"
            new_output_folder = f"{args.workspace_name}/output"
            if args.code: # is sent almost all the time
                new_log_folder = f"{args.workspace_name}/logs/{args.code}"
                new_output_folder = f"{args.workspace_name}/output/{args.code}"
                            
            if not os.path.exists(new_log_folder):
                os.makedirs(new_log_folder)
            if not os.path.exists(new_output_folder):
                os.makedirs(new_output_folder)
                
            print(f"Creating BERT Trainer .... masking: True, max_mask: {args.max_mask}")
            trainer = BERTTrainer(bert, len(vocab_obj.vocab), train_dataloader=pretrain_data_loader, 
                                  val_dataloader=pretrain_val_data_loader, test_dataloader=pretrain_test_data_loader,
                                  lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, 
                                  with_cuda=args.with_cuda, cuda_devices=args.cuda_devices, log_freq=args.log_freq, 
                                  log_folder_path=new_log_folder)

            start_time = time.time()
            print(f'Pretraining Starts, Time: {time.strftime("%D %T", time.localtime(start_time))}')
            # if need to pretrain from a check-point, need :check_epoch
            repoch = range(args.check_epoch, args.epochs) if args.check_epoch else range(args.epochs)
            counter = 0
            patience = 20
            for epoch in repoch:
                print(f'Training Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}')
                trainer.train(epoch)
                print(f'Training Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n')

                if pretrain_val_data_loader is not None:
                    print(f'Validation Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}')
                    trainer.val(epoch)
                    print(f'Validation Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n')

                if trainer.save_model: #  or epoch%10 == 0 and epoch > 4  
                    trainer.save(epoch, args.output_path)
                    counter = 0
                    if pretrain_test_data_loader is not None:
                        print(f'Test Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}')
                        trainer.test(epoch)
                        print(f'Test Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n')
                else:
                    counter +=1
                    if counter >= patience:
                        print(f"Early stopping at epoch {epoch}")
                        break

            end_time = time.time()
            print("Time Taken to pretrain model = ", end_time - start_time)
            print(f'Pretraining Ends, Time: {time.strftime("%D %T", time.localtime(end_time))}')
        else:
            print("Fine Tuning......")
            print("Loading Train Dataset", args.train_dataset_path)            
            train_dataset = TokenizerDataset(args.train_dataset_path, args.train_label_path, vocab_obj, seq_len=args.seq_len)

#             print("Loading Validation Dataset", args.val_dataset_path)            
#             val_dataset = TokenizerDataset(args.val_dataset_path, args.val_label_path, vocab_obj, seq_len=args.seq_len) \
#                 if args.val_dataset_path is not None else None
            
            print("Loading Test Dataset", args.test_dataset_path)
            test_dataset = TokenizerDataset(args.test_dataset_path, args.test_label_path, vocab_obj, seq_len=args.seq_len) \
                if args.test_dataset_path is not None else None

            print("Creating Dataloader...")
            train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)
            # val_data_loader = DataLoader(val_dataset, batch_size=args.batch_size, num_workers=args.num_workers) \
            #     if val_dataset is not None else None
            test_data_loader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers) \
                if test_dataset is not None else None

            print("Load Pre-trained BERT model")
            # bert = BERT(len(vocab_obj.vocab), hidden=args.hidden, n_layers=args.layers, attn_heads=args.attn_heads)
            cuda_condition = torch.cuda.is_available() and args.with_cuda
            device = torch.device("cuda:0" if cuda_condition else "cpu")
            bert = torch.load(args.pretrained_bert_checkpoint, map_location=device)

    #         if args.finetune_task == "SL":
    #             if args.workspace_name == "ratio_proportion_change4":
    #                 num_labels = 9
    #             elif args.workspace_name == "ratio_proportion_change3":
    #                 num_labels = 9
    #             elif args.workspace_name == "scale_drawings_3":
    #                 num_labels = 9
    #             elif args.workspace_name == "sales_tax_discounts_two_rates":
    #                 num_labels = 3
    #         else:
            # num_labels = 2
    #         # num_labels = 1
            # print(f"Number of Labels : {args.num_labels}")
            new_log_folder = f"{args.workspace_name}/logs"
            new_output_folder = f"{args.workspace_name}/output"
            if args.finetune_task: # is sent almost all the time
                new_log_folder = f"{args.workspace_name}/logs/{args.finetune_task}"
                new_output_folder = f"{args.workspace_name}/output/{args.finetune_task}"
                            
            if not os.path.exists(new_log_folder):
                os.makedirs(new_log_folder)
            if not os.path.exists(new_output_folder):
                os.makedirs(new_output_folder)
                
            print("Creating BERT Fine Tune Trainer")
            trainer = BERTFineTuneTrainer(bert, len(vocab_obj.vocab), 
                          train_dataloader=train_data_loader, test_dataloader=test_data_loader, 
                          lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, 
                          with_cuda=args.with_cuda, cuda_devices = args.cuda_devices, log_freq=args.log_freq, 
                          workspace_name = args.workspace_name, num_labels=args.num_labels, log_folder_path=new_log_folder)

            print("Fine-tune training Start....")
            start_time = time.time()
            repoch = range(args.check_epoch, args.epochs) if args.check_epoch else range(args.epochs)
            counter = 0
            patience = 10
            for epoch in repoch:
                print(f'Training Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}')
                trainer.train(epoch)
                print(f'Training Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n')

                if test_data_loader is not None:
                    print(f'Test Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}')
                    trainer.test(epoch)
                    # pickle.dump(trainer.probability_list, open(f"{args.workspace_name}/output/aaai/change4_mid_prob_{epoch}.pkl","wb"))
                    print(f'Test Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n')

                # if val_data_loader is not None:
                #     print(f'Validation Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}')
                #     trainer.val(epoch)
                #     print(f'Validation Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n')
                
                if trainer.save_model: #  or epoch%10 == 0
                    trainer.save(epoch, args.output_path)
                    counter = 0
                else:
                    counter +=1
                    if counter >= patience:
                        print(f"Early stopping at epoch {epoch}")
                        break
                    
            end_time = time.time()
            print("Time Taken to fine-tune model = ", end_time - start_time)
            print(f'Pretraining Ends, Time: {time.strftime("%D %T", time.localtime(end_time))}')
            

if __name__ == "__main__":
    train()