model/model.py

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()]