Upload model_manager.py

model_manager.py

@@ -0,0 +1,324 @@
+'''
+Author: Qiguang Chen
+Date: 2023-01-11 10:39:26
+LastEditors: Qiguang Chen
+LastEditTime: 2023-02-08 00:42:56
+Description: manage all process of model training and prediction.
+
+'''
+import os
+import random
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+
+from common import utils
+from common.loader import DataFactory
+from common.logger import Logger
+from common.metric import Evaluator
+from common.tokenizer import get_tokenizer, get_tokenizer_class, load_embedding
+from common.utils import InputData, instantiate
+from common.utils import OutputData
+from common.config import Config
+import dill
+
+
+class ModelManager(object):
+    def __init__(self, config: Config):
+        """create model manager by config
+
+        Args:
+            config (Config): configuration to manage all process in OpenSLU
+        """
+        # init config
+        self.config = config
+        self.__set_seed(self.config.base.get("seed"))
+        self.device = self.config.base.get("device")
+        
+        # enable accelerator
+        if "accelerator" in self.config and self.config["accelerator"].get("use_accelerator"):
+            from accelerate import Accelerator
+            self.accelerator = Accelerator(log_with="wandb")
+        else:
+            self.accelerator = None
+        if self.config.base.get("train"):
+            self.tokenizer = get_tokenizer(
+                self.config.tokenizer.get("_tokenizer_name_"))
+            self.logger = Logger(
+                "wandb", self.config.base["name"], start_time=config.start_time, accelerator=self.accelerator)
+
+            # init dataloader & load data
+            if self.config.base.get("save_dir"):
+                self.model_save_dir = self.config.base["save_dir"]
+            else:
+                if not os.path.exists("save/"):
+                    os.mkdir("save/")
+                self.model_save_dir = "save/" + config.start_time
+            if not os.path.exists(self.model_save_dir):
+                    os.mkdir(self.model_save_dir)
+            batch_size = self.config.base["batch_size"]
+            df = DataFactory(tokenizer=self.tokenizer,
+                             use_multi_intent=self.config.base.get("multi_intent"),
+                             to_lower_case=self.config.base.get("_to_lower_case_"))
+            train_dataset = df.load_dataset(self.config.dataset, split="train")
+
+            # update label and vocabulary
+            df.update_label_names(train_dataset)
+            df.update_vocabulary(train_dataset)
+
+            # init tokenizer config and dataloaders
+            tokenizer_config = {key: self.config.tokenizer[key]
+                                for key in self.config.tokenizer if key[0] != "_" and key[-1] != "_"}
+            self.train_dataloader = df.get_data_loader(train_dataset,
+                                                       batch_size,
+                                                       shuffle=True,
+                                                       device=self.device,
+                                                       enable_label=True,
+                                                       align_mode=self.config.tokenizer.get(
+                                                           "_align_mode_"),
+                                                       label2tensor=True,
+                                                       **tokenizer_config)
+            dev_dataset = df.load_dataset(
+                self.config.dataset, split="validation")
+            self.dev_dataloader = df.get_data_loader(dev_dataset,
+                                                     batch_size,
+                                                     shuffle=False,
+                                                     device=self.device,
+                                                     enable_label=True,
+                                                     align_mode=self.config.tokenizer.get(
+                                                         "_align_mode_"),
+                                                     label2tensor=False,
+                                                     **tokenizer_config)
+            df.update_vocabulary(dev_dataset)
+            # add intent label num and slot label num to config
+            if int(self.config.get_intent_label_num()) == 0 or int(self.config.get_slot_label_num()) == 0:
+                self.intent_list = df.intent_label_list
+                self.intent_dict = df.intent_label_dict
+                self.config.set_intent_label_num(len(self.intent_list))
+                self.slot_list = df.slot_label_list
+                self.slot_dict = df.slot_label_dict
+                self.config.set_slot_label_num(len(self.slot_list))
+                self.config.set_vocab_size(self.tokenizer.vocab_size)
+
+            # autoload embedding for non-pretrained encoder
+            if self.config["model"]["encoder"].get("embedding") and self.config["model"]["encoder"]["embedding"].get(
+                    "load_embedding_name"):
+                self.config["model"]["encoder"]["embedding"]["embedding_matrix"] = load_embedding(self.tokenizer,
+                                                                                                  self.config["model"][
+                                                                                                      "encoder"][
+                                                                                                      "embedding"].get(
+                                                                                                      "load_embedding_name"))
+            # fill template in config
+            self.config.autoload_template()
+            # save config
+            self.logger.set_config(self.config)
+
+            self.model = None
+            self.optimizer = None
+            self.total_step = None
+            self.lr_scheduler = None
+            if self.config.tokenizer.get("_tokenizer_name_") == "word_tokenizer":
+                self.tokenizer.save(os.path.join(self.model_save_dir, "tokenizer.json"))
+            utils.save_json(os.path.join(
+                self.model_save_dir, "label.json"), {"intent": self.intent_list,"slot": self.slot_list})
+        if self.config.base.get("test"):
+            self.test_dataset = df.load_dataset(
+                self.config.dataset, split="test")
+            self.test_dataloader = df.get_data_loader(self.test_dataset,
+                                                      batch_size,
+                                                      shuffle=False,
+                                                      device=self.device,
+                                                      enable_label=True,
+                                                      align_mode=self.config.tokenizer.get(
+                                                          "_align_mode_"),
+                                                      label2tensor=False,
+                                                      **tokenizer_config)
+
+    def init_model(self, model):
+        """init model, optimizer, lr_scheduler
+
+        Args:
+            model (Any): pytorch model
+        """
+        self.model = model
+        self.model.to(self.device)
+        if self.config.base.get("train"):
+            self.optimizer = instantiate(
+                self.config["optimizer"])(self.model.parameters())
+            self.total_step = int(self.config.base.get(
+                "epoch_num")) * len(self.train_dataloader)
+            self.lr_scheduler = instantiate(self.config["scheduler"])(
+                optimizer=self.optimizer,
+                num_training_steps=self.total_step
+            )
+        if self.accelerator is not None:
+            self.model, self.optimizer, self.train_dataloader, self.lr_scheduler = self.accelerator.prepare(
+                self.model, self.optimizer, self.train_dataloader, self.lr_scheduler)
+            if self.config.base.get("load_dir_path"):
+                self.accelerator.load_state(self.config.base.get("load_dir_path"))
+            # self.dev_dataloader = self.accelerator.prepare(self.dev_dataloader)
+
+    def eval(self, step: int, best_metric: float) -> float:
+        """ evaluation models.
+
+        Args:
+            step (int): which step the model has trained in
+            best_metric (float): last best metric value to judge whether to test or save model
+
+        Returns:
+            float: updated best metric value
+        """
+        # TODO: save dev
+        _, res = self.__evaluate(self.model, self.dev_dataloader)
+        self.logger.log_metric(res, metric_split="dev", step=step)
+        if res[self.config.base.get("best_key")] > best_metric:
+            best_metric = res[self.config.base.get("best_key")]
+            outputs, test_res = self.__evaluate(
+                self.model, self.test_dataloader)
+            if not os.path.exists(self.model_save_dir):
+                os.mkdir(self.model_save_dir)
+            if self.accelerator is None:
+                torch.save(self.model, os.path.join(
+                    self.model_save_dir, "model.pkl"))
+                torch.save(self.optimizer, os.path.join(
+                    self.model_save_dir, "optimizer.pkl"))
+                torch.save(self.lr_scheduler, os.path.join(
+                    self.model_save_dir, "lr_scheduler.pkl"), pickle_module=dill)
+                torch.save(step, os.path.join(
+                    self.model_save_dir, "step.pkl"))
+            else:
+                self.accelerator.wait_for_everyone()
+                unwrapped_model = self.accelerator.unwrap_model(self.model)
+                self.accelerator.save(unwrapped_model.state_dict(
+                ), os.path.join(self.model_save_dir, "model.pkl"))
+                self.accelerator.save_state(output_dir=self.model_save_dir)
+            outputs.save(self.model_save_dir, self.test_dataset)
+            self.logger.log_metric(test_res, metric_split="test", step=step)
+        return best_metric
+
+    def train(self) -> float:
+        """ train models.
+
+        Returns:
+            float: updated best metric value
+        """
+        step = 0
+        best_metric = 0
+        progress_bar = tqdm(range(self.total_step))
+        for _ in range(int(self.config.base.get("epoch_num"))):
+            for data in self.train_dataloader:
+                if step == 0:
+                    self.logger.info(data.get_item(
+                        0, tokenizer=self.tokenizer, intent_map=self.intent_list, slot_map=self.slot_list))
+                output = self.model(data)
+                if self.accelerator is not None and hasattr(self.model, "module"):
+                    loss, intent_loss, slot_loss = self.model.module.compute_loss(
+                        pred=output, target=data)
+                else:
+                    loss, intent_loss, slot_loss = self.model.compute_loss(
+                        pred=output, target=data)
+                self.logger.log_loss(loss, "Loss", step=step)
+                self.logger.log_loss(intent_loss, "Intent Loss", step=step)
+                self.logger.log_loss(slot_loss, "Slot Loss", step=step)
+                self.optimizer.zero_grad()
+
+                if self.accelerator is not None:
+                    self.accelerator.backward(loss)
+                else:
+                    loss.backward()
+                self.optimizer.step()
+                self.lr_scheduler.step()
+                if not self.config.base.get("eval_by_epoch") and step % self.config.base.get(
+                        "eval_step") == 0 and step != 0:
+                    best_metric = self.eval(step, best_metric)
+                step += 1
+                progress_bar.update(1)
+            if self.config.base.get("eval_by_epoch"):
+                best_metric = self.eval(step, best_metric)
+        self.logger.finish()
+        return best_metric
+
+    def __set_seed(self, seed_value: int):
+        """Manually set random seeds.
+
+        Args:
+            seed_value (int): random seed
+        """
+        random.seed(seed_value)
+        np.random.seed(seed_value)
+        torch.manual_seed(seed_value)
+        torch.random.manual_seed(seed_value)
+        os.environ['PYTHONHASHSEED'] = str(seed_value)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed_value)
+            torch.cuda.manual_seed_all(seed_value)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = True
+        return
+
+    def __evaluate(self, model, dataloader):
+        model.eval()
+        inps = InputData()
+        outputs = OutputData()
+        for data in dataloader:
+            torch.cuda.empty_cache()
+            output = model(data)
+            if self.accelerator is not None and hasattr(self.model, "module"):
+                decode_output = model.module.decode(output, data)
+            else:
+                decode_output = model.decode(output, data)
+
+            decode_output.map_output(slot_map=self.slot_list,
+                                     intent_map=self.intent_list)
+            data, decode_output = utils.remove_slot_ignore_index(
+                data, decode_output, ignore_index="#")
+
+            inps.merge_input_data(data)
+            outputs.merge_output_data(decode_output)
+        if "metric" in self.config:
+            res = Evaluator.compute_all_metric(
+                inps, outputs, intent_label_map=self.intent_dict, metric_list=self.config.metric)
+        else:
+            res = Evaluator.compute_all_metric(
+                inps, outputs, intent_label_map=self.intent_dict)
+        model.train()
+        return outputs, res
+
+    def load(self):
+        
+        self.model = torch.load(os.path.join(self.config.base["model_dir"], "model.pkl"),map_location=self.config.base["device"])
+        if self.config.tokenizer["_tokenizer_name_"] == "word_tokenizer":
+            self.tokenizer = get_tokenizer_class(self.config.tokenizer["_tokenizer_name_"]).from_file(
+                os.path.join(self.config.base["model_dir"], "tokenizer.json"))
+        else:
+            self.tokenizer = get_tokenizer(self.config.tokenizer["_tokenizer_name_"])
+        self.model.to(self.device)
+        label = utils.load_json(os.path.join(self.config.base["model_dir"], "label.json"))
+        self.intent_list = label["intent"]
+        self.slot_list = label["slot"]
+        self.data_factory=DataFactory(tokenizer=self.tokenizer,
+                             use_multi_intent=self.config.base.get("multi_intent"),
+                             to_lower_case=self.config.tokenizer.get("_to_lower_case_"))
+
+    def predict(self, text_data):
+        self.model.eval()
+        tokenizer_config = {key: self.config.tokenizer[key]
+                                for key in self.config.tokenizer if key[0] != "_" and key[-1] != "_"}
+        align_mode = self.config.tokenizer.get("_align_mode_")
+        inputs = self.data_factory.batch_fn(batch=[{"text": text_data.split(" ")}],
+                                          device=self.device,
+                                        config=tokenizer_config,
+                                        enable_label=False,
+                                        align_mode= align_mode if align_mode is not None else "general",
+                                        label2tensor=False)
+        output = self.model(inputs)
+        decode_output = self.model.decode(output, inputs)
+        decode_output.map_output(slot_map=self.slot_list,
+                                     intent_map=self.intent_list)
+        if self.config.base.get("multi_intent"):
+            intent = decode_output.intent_ids[0]
+        else:
+            intent = [decode_output.intent_ids[0]]
+        return {"intent": intent, "slot": decode_output.slot_ids[0], "text": self.tokenizer.decode(inputs.input_ids[0])}