Delete model_manager.py

model_manager.py

@@ -1,324 +0,0 @@
-'''
-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])}