model_seq

model/model.py

@@ -0,0 +1,511 @@
+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()]