Update model/model.py

model/model.py

@@ -1,511 +1,517 @@
-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):
+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
+
+MODEL_CLASSES = {
+    "auto": (AutoConfig, AutoModel, AutoTokenizer),
+    #"mbart": (MBartConfig, MBartForConditionalGeneration, MBartTokenizer),
+    "bartpho": (MBartConfig, MBartForConditionalGeneration, AutoTokenizer)
+}
+
+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()]