from bartpho.preprocess import normalize, tokenize from bartpho.utils import tag_dict, polarity_dict, polarity_list, tags, eng_tags, eng_polarity, detect_labels, no_polarity, no_tag from bartpho.utils import predict, predict_df, predict_detect, predict_df_detect from simpletransformers.config.model_args import Seq2SeqArgs import random import numpy as np import torch from transformers import ( AdamW, AutoConfig, AutoModel, AutoTokenizer, MBartConfig, MBartForConditionalGeneration, MBartTokenizer, get_linear_schedule_with_warmup, ) from pyvi.ViTokenizer import tokenize as model_tokenize class Seq2SeqModel: def __init__( self, encoder_decoder_type=None, encoder_decoder_name=None, config=None, args=None, use_cuda=False, cuda_device=0, **kwargs, ): """ Initializes a Seq2SeqModel. Args: encoder_decoder_type (optional): The type of encoder-decoder model. (E.g. bart) encoder_decoder_name (optional): The path to a directory containing the saved encoder and decoder of a Seq2SeqModel. (E.g. "outputs/") OR a valid BART or MarianMT model. config (optional): A configuration file to build an EncoderDecoderModel. args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args. use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only. cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default. **kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied. """ # noqa: ignore flake8" if not config: # if not ((encoder_name and decoder_name) or encoder_decoder_name) and not encoder_type: if not encoder_decoder_name: raise ValueError( "You must specify a Seq2Seq config \t OR \t" "encoder_decoder_name" ) elif not encoder_decoder_type: raise ValueError( "You must specify a Seq2Seq config \t OR \t" "encoder_decoder_name" ) self.args = self._load_model_args(encoder_decoder_name) print(args) if args: self.args.update_from_dict(args) print(args) if self.args.manual_seed: random.seed(self.args.manual_seed) np.random.seed(self.args.manual_seed) torch.manual_seed(self.args.manual_seed) if self.args.n_gpu > 0: torch.cuda.manual_seed_all(self.args.manual_seed) if use_cuda: if torch.cuda.is_available(): self.device = torch.device("cuda") else: raise ValueError( "'use_cuda' set to True when cuda is unavailable." "Make sure CUDA is available or set `use_cuda=False`." ) else: self.device = "cpu" self.results = {} if not use_cuda: self.args.fp16 = False # config = EncoderDecoderConfig.from_encoder_decoder_configs(config, config) #if encoder_decoder_type: config_class, model_class, tokenizer_class = MODEL_CLASSES[encoder_decoder_type] self.model = model_class.from_pretrained(encoder_decoder_name) self.encoder_tokenizer = tokenizer_class.from_pretrained(encoder_decoder_name) self.decoder_tokenizer = self.encoder_tokenizer self.config = self.model.config if self.args.wandb_project and not wandb_available: warnings.warn("wandb_project specified but wandb is not available. Wandb disabled.") self.args.wandb_project = None self.args.model_name = encoder_decoder_name self.args.model_type = encoder_decoder_type def train_model( self, train_data, best_accuracy, output_dir=None, show_running_loss=True, args=None, eval_data=None, test_data=None, verbose=True, **kwargs, ): if args: self.args.update_from_dict(args) #self.args = args if self.args.silent: show_running_loss = False if not output_dir: output_dir = self.args.output_dir self._move_model_to_device() train_dataset = self.load_and_cache_examples(train_data, verbose=verbose) os.makedirs(output_dir, exist_ok=True) global_step, tr_loss, best_accuracy = self.train( train_dataset, output_dir, best_accuracy, show_running_loss=show_running_loss, eval_data=eval_data, test_data=test_data, verbose=verbose, **kwargs, ) final_dir = self.args.output_dir + "/final" self._save_model(final_dir, model=self.model) if verbose: logger.info(" Training of {} model complete. Saved best to {}.".format(self.args.model_name, final_dir)) return best_accuracy def train( self, train_dataset, output_dir, best_accuracy, show_running_loss=True, eval_data=None, test_data=None, verbose=True, **kwargs, ): """ Trains the model on train_dataset. Utility function to be used by the train_model() method. Not intended to be used directly. """ #epoch_lst = [] #acc_detects, pre_detects, rec_detects, f1_detects, accs, pre_absas, rec_absas, f1_absas = [], [], [], [], [], [], [], [] #tacc_detects, tpre_detects, trec_detects, tf1_detects, taccs, tpre_absas, trec_absas, tf1_absas = [], [], [], [], [], [], [], [] model = self.model args = self.args tb_writer = SummaryWriter(logdir=args.tensorboard_dir) train_sampler = RandomSampler(train_dataset) train_dataloader = DataLoader( train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, num_workers=self.args.dataloader_num_workers, ) if args.max_steps > 0: t_total = args.max_steps args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1 else: t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs no_decay = ["bias", "LayerNorm.weight"] optimizer_grouped_parameters = [] custom_parameter_names = set() for group in self.args.custom_parameter_groups: params = group.pop("params") custom_parameter_names.update(params) param_group = {**group} param_group["params"] = [p for n, p in model.named_parameters() if n in params] optimizer_grouped_parameters.append(param_group) for group in self.args.custom_layer_parameters: layer_number = group.pop("layer") layer = f"layer.{layer_number}." group_d = {**group} group_nd = {**group} group_nd["weight_decay"] = 0.0 params_d = [] params_nd = [] for n, p in model.named_parameters(): if n not in custom_parameter_names and layer in n: if any(nd in n for nd in no_decay): params_nd.append(p) else: params_d.append(p) custom_parameter_names.add(n) group_d["params"] = params_d group_nd["params"] = params_nd optimizer_grouped_parameters.append(group_d) optimizer_grouped_parameters.append(group_nd) if not self.args.train_custom_parameters_only: optimizer_grouped_parameters.extend( [ { "params": [ p for n, p in model.named_parameters() if n not in custom_parameter_names and not any(nd in n for nd in no_decay) ], "weight_decay": args.weight_decay, }, { "params": [ p for n, p in model.named_parameters() if n not in custom_parameter_names and any(nd in n for nd in no_decay) ], "weight_decay": 0.0, }, ] ) warmup_steps = math.ceil(t_total * args.warmup_ratio) args.warmup_steps = warmup_steps if args.warmup_steps == 0 else args.warmup_steps # TODO: Use custom optimizer like with BertSum? optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total ) if (args.model_name and os.path.isfile(os.path.join(args.model_name, "optimizer.pt")) and os.path.isfile(os.path.join(args.model_name, "scheduler.pt"))): # Load in optimizer and scheduler states optimizer.load_state_dict(torch.load(os.path.join(args.model_name, "optimizer.pt"))) scheduler.load_state_dict(torch.load(os.path.join(args.model_name, "scheduler.pt"))) if args.n_gpu > 1: model = torch.nn.DataParallel(model) logger.info(" Training started") global_step = 0 tr_loss, logging_loss = 0.0, 0.0 model.zero_grad() train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.silent, mininterval=0) epoch_number = 0 best_eval_metric = None early_stopping_counter = 0 steps_trained_in_current_epoch = 0 epochs_trained = 0 if args.model_name and os.path.exists(args.model_name): try: # set global_step to gobal_step of last saved checkpoint from model path checkpoint_suffix = args.model_name.split("/")[-1].split("-") if len(checkpoint_suffix) > 2: checkpoint_suffix = checkpoint_suffix[1] else: checkpoint_suffix = checkpoint_suffix[-1] global_step = int(checkpoint_suffix) epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps) steps_trained_in_current_epoch = global_step % ( len(train_dataloader) // args.gradient_accumulation_steps ) logger.info(" Continuing training from checkpoint, will skip to saved global_step") logger.info(" Continuing training from epoch %d", epochs_trained) logger.info(" Continuing training from global step %d", global_step) logger.info(" Will skip the first %d steps in the current epoch", steps_trained_in_current_epoch) except ValueError: logger.info(" Starting fine-tuning.") if args.wandb_project: wandb.init(project=args.wandb_project, config={**asdict(args)}, **args.wandb_kwargs) wandb.watch(self.model) if args.fp16: from torch.cuda import amp scaler = amp.GradScaler() model.train() for current_epoch in train_iterator: if epochs_trained > 0: epochs_trained -= 1 continue train_iterator.set_description(f"Epoch {epoch_number + 1} of {args.num_train_epochs}") batch_iterator = tqdm( train_dataloader, desc=f"Running Epoch {epoch_number} of {args.num_train_epochs}", disable=args.silent, mininterval=0, ) for step, batch in enumerate(batch_iterator): if steps_trained_in_current_epoch > 0: steps_trained_in_current_epoch -= 1 continue # batch = tuple(t.to(device) for t in batch) inputs = self._get_inputs_dict(batch) if args.fp16: with amp.autocast(): outputs = model(**inputs) # model outputs are always tuple in pytorch-transformers (see doc) loss = outputs[0] else: outputs = model(**inputs) # model outputs are always tuple in pytorch-transformers (see doc) loss = outputs[0] if args.n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu parallel training current_loss = loss.item() if show_running_loss: batch_iterator.set_description( f"Epochs {epoch_number}/{args.num_train_epochs}. Running Loss: {current_loss:9.4f}" ) if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps if args.fp16: scaler.scale(loss).backward() else: loss.backward() tr_loss += loss.item() if (step + 1) % args.gradient_accumulation_steps == 0: if args.fp16: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) if args.fp16: scaler.step(optimizer) scaler.update() else: optimizer.step() scheduler.step() # Update learning rate schedule model.zero_grad() global_step += 1 if args.logging_steps > 0 and global_step % args.logging_steps == 0: # Log metrics tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step) tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step) logging_loss = tr_loss if args.wandb_project: wandb.log( { "Training loss": current_loss, "lr": scheduler.get_lr()[0], "global_step": global_step, } ) # if args.save_steps > 0 and global_step % args.save_steps == 0: # # Save model checkpoint # output_dir_current = os.path.join(output_dir, "checkpoint-{}".format(global_step)) # self._save_model(output_dir_current, optimizer, scheduler, model=model) epoch_number += 1 output_dir_current = os.path.join(output_dir, "checkpoint-{}-epoch-{}".format(global_step, epoch_number)) print('batch: '+str(args.train_batch_size)+' accumulation_steps: '+str(args.gradient_accumulation_steps)+\ ' lr: '+str(args.learning_rate)+' epochs: '+str(args.num_train_epochs)+' epoch: '+str(epoch_number)) print('---dev dataset----') acc_detect, pre_detect, rec_detect, f1_detect, acc, pre_absa, rec_absa, f1_absa = predict_df(model, eval_data, tokenizer=self.encoder_tokenizer, device=self.device) print('---test dataset----') tacc_detect, tpre_detect, trec_detect, tf1_detect, tacc, tpre_absa, trec_absa, tf1_absa = predict_df(model, test_data, tokenizer=self.encoder_tokenizer, device=self.device) # if acc > best_accuracy: # best_accuracy = acc # if not args.save_model_every_epoch: # self._save_model(output_dir_current, optimizer, scheduler, model=model) # with open('./MAMS_best_accuracy.txt', 'a') as f0: # f0.writelines('batch: '+str(args.train_batch_size)+' accumulation_steps: '+str(args.gradient_accumulation_steps)+\ # ' lr: '+str(args.learning_rate)+' epochs: '+str(args.num_train_epochs)+' epoch: '+str(epoch_number)+' val_accuracy: '+str(best_accuracy)+\ # ' test_accuracy: '+str(tacc)+'\n') # if args.save_model_every_epoch: # os.makedirs(output_dir_current, exist_ok=True) # self._save_model(output_dir_current, optimizer, scheduler, model=model) if acc > best_accuracy: # Cập nhật best_accuracy nếu tìm thấy mô hình tốt hơn best_accuracy = acc # Lưu mô hình tốt nhất vào output_dir_current self._save_model(output_dir_current, optimizer, scheduler, model=model) # Ghi lại thông tin về best_accuracy vào file log with open('./MAMS_best_accuracy.txt', 'a') as f0: f0.writelines( 'batch: ' + str(args.train_batch_size) + ' accumulation_steps: ' + str(args.gradient_accumulation_steps) + ' lr: ' + str(args.learning_rate) + ' epochs: ' + str(args.num_train_epochs) + ' epoch: ' + str(epoch_number) + ' val_accuracy: ' + str(best_accuracy) + ' test_accuracy: ' + str(tacc) + '\n' ) return global_step, tr_loss / global_step, best_accuracy def load_and_cache_examples(self, data, evaluate=False, no_cache=False, verbose=True, silent=False): """ Creates a T5Dataset from data. Utility function for train() and eval() methods. Not intended to be used directly. """ encoder_tokenizer = self.encoder_tokenizer decoder_tokenizer = self.decoder_tokenizer args = self.args if not no_cache: no_cache = args.no_cache if not no_cache: os.makedirs(self.args.cache_dir, exist_ok=True) mode = "dev" if evaluate else "train" if args.dataset_class: CustomDataset = args.dataset_class return CustomDataset(encoder_tokenizer, decoder_tokenizer, args, data, mode) else: return SimpleSummarizationDataset(encoder_tokenizer, self.args, data, mode) def _save_model(self, output_dir=None, optimizer=None, scheduler=None, model=None, results=None): if not output_dir: output_dir = self.args.output_dir os.makedirs(output_dir, exist_ok=True) logger.info(f"Saving model into {output_dir}") if model and not self.args.no_save: # Take care of distributed/parallel training model_to_save = model.module if hasattr(model, "module") else model self._save_model_args(output_dir) os.makedirs(os.path.join(output_dir), exist_ok=True) model_to_save.save_pretrained(output_dir) self.config.save_pretrained(output_dir) self.encoder_tokenizer.save_pretrained(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) if optimizer and scheduler and self.args.save_optimizer_and_scheduler: torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt")) torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt")) if results: output_eval_file = os.path.join(output_dir, "eval_results.txt") with open(output_eval_file, "w") as writer: for key in sorted(results.keys()): writer.write("{} = {}\n".format(key, str(results[key]))) def _move_model_to_device(self): self.model.to(self.device) def _get_inputs_dict(self, batch): device = self.device pad_token_id = self.encoder_tokenizer.pad_token_id source_ids, source_mask, y = batch["source_ids"], batch["source_mask"], batch["target_ids"] y_ids = y[:, :-1].contiguous() lm_labels = y[:, 1:].clone() lm_labels[y[:, 1:] == pad_token_id] = -100 inputs = { "input_ids": source_ids.to(device), "attention_mask": source_mask.to(device), "decoder_input_ids": y_ids.to(device), "labels": lm_labels.to(device), } return inputs def _save_model_args(self, output_dir): os.makedirs(output_dir, exist_ok=True) self.args.save(output_dir) def _load_model_args(self, input_dir): args = Seq2SeqArgs() args.load(input_dir) return args def get_named_parameters(self): return [n for n, p in self.model.named_parameters()]