Spaces:
Running
Running
| 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() |