import logging import os import numpy as np import torch from early_stopping import EarlyStopping from torch.utils.data import DataLoader, RandomSampler, SequentialSampler from torch.utils.tensorboard import SummaryWriter from tqdm.auto import tqdm, trange from transformers import AdamW, get_linear_schedule_with_warmup from utils import MODEL_CLASSES, compute_metrics, get_intent_labels, get_slot_labels logger = logging.getLogger(__name__) class Trainer(object): def __init__(self, args, train_dataset=None, dev_dataset=None, test_dataset=None): self.args = args self.train_dataset = train_dataset self.dev_dataset = dev_dataset self.test_dataset = test_dataset self.intent_label_lst = get_intent_labels(args) self.slot_label_lst = get_slot_labels(args) # Use cross entropy ignore index as padding label id so that only real label ids contribute to the loss later self.pad_token_label_id = args.ignore_index self.config_class, self.model_class, _ = MODEL_CLASSES[args.model_type] # self.config = self.config_class.from_pretrained(model_path, finetuning_task=args.task) if args.pretrained: print(args.model_name_or_path) self.model = self.model_class.from_pretrained( args.pretrained_path, args=args, intent_label_lst=self.intent_label_lst, slot_label_lst=self.slot_label_lst, ) else: self.config = self.config_class.from_pretrained(args.model_name_or_path, finetuning_task=args.token_level) self.model = self.model_class.from_pretrained( args.model_name_or_path, config=self.config, args=args, intent_label_lst=self.intent_label_lst, slot_label_lst=self.slot_label_lst, ) # GPU or CPU torch.cuda.set_device(self.args.gpu_id) print(self.args.gpu_id) print(torch.cuda.current_device()) self.device = "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu" self.model.to(self.device) def train(self): train_sampler = RandomSampler(self.train_dataset) train_dataloader = DataLoader(self.train_dataset, sampler=train_sampler, batch_size=self.args.train_batch_size) writer = SummaryWriter(log_dir=self.args.model_dir) if self.args.max_steps > 0: t_total = self.args.max_steps self.args.num_train_epochs = ( self.args.max_steps // (len(train_dataloader) // self.args.gradient_accumulation_steps) + 1 ) else: t_total = len(train_dataloader) // self.args.gradient_accumulation_steps * self.args.num_train_epochs print("check init") results = self.evaluate("dev") print(results) # Prepare optimizer and schedule (linear warmup and decay) no_decay = ["bias", "LayerNorm.weight"] optimizer_grouped_parameters = [ { "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": self.args.weight_decay, }, { "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0, }, ] optimizer = AdamW(optimizer_grouped_parameters, lr=self.args.learning_rate, eps=self.args.adam_epsilon) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=self.args.warmup_steps, num_training_steps=t_total ) # Train! logger.info("***** Running training *****") logger.info(" Num examples = %d", len(self.train_dataset)) logger.info(" Num Epochs = %d", self.args.num_train_epochs) logger.info(" Total train batch size = %d", self.args.train_batch_size) logger.info(" Gradient Accumulation steps = %d", self.args.gradient_accumulation_steps) logger.info(" Total optimization steps = %d", t_total) logger.info(" Logging steps = %d", self.args.logging_steps) logger.info(" Save steps = %d", self.args.save_steps) global_step = 0 tr_loss = 0.0 self.model.zero_grad() train_iterator = trange(int(self.args.num_train_epochs), desc="Epoch") early_stopping = EarlyStopping(patience=self.args.early_stopping, verbose=True) for _ in train_iterator: epoch_iterator = tqdm(train_dataloader, desc="Iteration", position=0, leave=True) print("\nEpoch", _) for step, batch in enumerate(epoch_iterator): self.model.train() batch = tuple(t.to(self.device) for t in batch) # GPU or CPU inputs = { "input_ids": batch[0], "attention_mask": batch[1], "intent_label_ids": batch[3], "slot_labels_ids": batch[4], } if self.args.model_type != "distilbert": inputs["token_type_ids"] = batch[2] outputs = self.model(**inputs) loss = outputs[0] if self.args.gradient_accumulation_steps > 1: loss = loss / self.args.gradient_accumulation_steps loss.backward() tr_loss += loss.item() if (step + 1) % self.args.gradient_accumulation_steps == 0: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm) optimizer.step() scheduler.step() # Update learning rate schedule self.model.zero_grad() global_step += 1 if self.args.logging_steps > 0 and global_step % self.args.logging_steps == 0: print("\nTuning metrics:", self.args.tuning_metric) results = self.evaluate("dev") writer.add_scalar("Loss/validation", results["loss"], _) writer.add_scalar("Intent Accuracy/validation", results["intent_acc"], _) writer.add_scalar("Slot F1/validation", results["slot_f1"], _) writer.add_scalar("Mean Intent Slot", results["mean_intent_slot"], _) writer.add_scalar("Sentence Accuracy/validation", results["semantic_frame_acc"], _) early_stopping(results[self.args.tuning_metric], self.model, self.args) if early_stopping.early_stop: print("Early stopping") break # if self.args.save_steps > 0 and global_step % self.args.save_steps == 0: # self.save_model() if 0 < self.args.max_steps < global_step: epoch_iterator.close() break if 0 < self.args.max_steps < global_step or early_stopping.early_stop: train_iterator.close() break writer.add_scalar("Loss/train", tr_loss / global_step, _) return global_step, tr_loss / global_step def write_evaluation_result(self, out_file, results): out_file = self.args.model_dir + "/" + out_file w = open(out_file, "w", encoding="utf-8") w.write("***** Eval results *****\n") for key in sorted(results.keys()): to_write = " {key} = {value}".format(key=key, value=str(results[key])) w.write(to_write) w.write("\n") w.close() def evaluate(self, mode): if mode == "test": dataset = self.test_dataset elif mode == "dev": dataset = self.dev_dataset else: raise Exception("Only dev and test dataset available") eval_sampler = SequentialSampler(dataset) eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=self.args.eval_batch_size) # Eval! logger.info("***** Running evaluation on %s dataset *****", mode) logger.info(" Num examples = %d", len(dataset)) logger.info(" Batch size = %d", self.args.eval_batch_size) eval_loss = 0.0 nb_eval_steps = 0 intent_preds = None slot_preds = None out_intent_label_ids = None out_slot_labels_ids = None self.model.eval() for batch in tqdm(eval_dataloader, desc="Evaluating"): batch = tuple(t.to(self.device) for t in batch) with torch.no_grad(): inputs = { "input_ids": batch[0], "attention_mask": batch[1], "intent_label_ids": batch[3], "slot_labels_ids": batch[4], } if self.args.model_type != "distilbert": inputs["token_type_ids"] = batch[2] outputs = self.model(**inputs) tmp_eval_loss, (intent_logits, slot_logits) = outputs[:2] eval_loss += tmp_eval_loss.mean().item() nb_eval_steps += 1 # Intent prediction if intent_preds is None: intent_preds = intent_logits.detach().cpu().numpy() out_intent_label_ids = inputs["intent_label_ids"].detach().cpu().numpy() else: intent_preds = np.append(intent_preds, intent_logits.detach().cpu().numpy(), axis=0) out_intent_label_ids = np.append( out_intent_label_ids, inputs["intent_label_ids"].detach().cpu().numpy(), axis=0 ) # Slot prediction if slot_preds is None: if self.args.use_crf: # decode() in `torchcrf` returns list with best index directly slot_preds = np.array(self.model.crf.decode(slot_logits)) else: slot_preds = slot_logits.detach().cpu().numpy() out_slot_labels_ids = inputs["slot_labels_ids"].detach().cpu().numpy() else: if self.args.use_crf: slot_preds = np.append(slot_preds, np.array(self.model.crf.decode(slot_logits)), axis=0) else: slot_preds = np.append(slot_preds, slot_logits.detach().cpu().numpy(), axis=0) out_slot_labels_ids = np.append( out_slot_labels_ids, inputs["slot_labels_ids"].detach().cpu().numpy(), axis=0 ) eval_loss = eval_loss / nb_eval_steps results = {"loss": eval_loss} # Intent result intent_preds = np.argmax(intent_preds, axis=1) # Slot result if not self.args.use_crf: slot_preds = np.argmax(slot_preds, axis=2) slot_label_map = {i: label for i, label in enumerate(self.slot_label_lst)} out_slot_label_list = [[] for _ in range(out_slot_labels_ids.shape[0])] slot_preds_list = [[] for _ in range(out_slot_labels_ids.shape[0])] for i in range(out_slot_labels_ids.shape[0]): for j in range(out_slot_labels_ids.shape[1]): if out_slot_labels_ids[i, j] != self.pad_token_label_id: out_slot_label_list[i].append(slot_label_map[out_slot_labels_ids[i][j]]) slot_preds_list[i].append(slot_label_map[slot_preds[i][j]]) total_result = compute_metrics(intent_preds, out_intent_label_ids, slot_preds_list, out_slot_label_list) results.update(total_result) logger.info("***** Eval results *****") for key in sorted(results.keys()): logger.info(" %s = %s", key, str(results[key])) if mode == "test": self.write_evaluation_result("eval_test_results.txt", results) elif mode == "dev": self.write_evaluation_result("eval_dev_results.txt", results) return results def save_model(self): # Save model checkpoint (Overwrite) if not os.path.exists(self.args.model_dir): os.makedirs(self.args.model_dir) model_to_save = self.model.module if hasattr(self.model, "module") else self.model model_to_save.save_pretrained(self.args.model_dir) # Save training arguments together with the trained model torch.save(self.args, os.path.join(self.args.model_dir, "training_args.bin")) logger.info("Saving model checkpoint to %s", self.args.model_dir) def load_model(self): # Check whether model exists if not os.path.exists(self.args.model_dir): raise Exception("Model doesn't exists! Train first!") try: self.model = self.model_class.from_pretrained( self.args.model_dir, args=self.args, intent_label_lst=self.intent_label_lst, slot_label_lst=self.slot_label_lst, ) self.model.to(self.device) logger.info("***** Model Loaded *****") except Exception: raise Exception("Some model files might be missing...")