configuration.py

import logging
import os

import configargparse


def init_logger(
    root_log_level=logging.DEBUG,
    console_log_level=logging.NOTSET,
    log_file=None,
    log_file_level=logging.NOTSET,
):
    """This funtion initializes a customized logger

    Keyword Arguments:
        root_log_level {int} -- root logging level (default: {logging.DEBUG})
        console_log_level {int} -- console logging level (default: {logging.NOTSET})
        log_file {str} -- logging file path (default: {None})
        log_file_level {int} -- logging file level (default: {logging.NOTSET})
    """

    log_format = logging.Formatter(
        "[%(asctime)s - %(filename)s - line:%(lineno)d - %(levelname)s]: %(message)s"
    )
    handlers = []

    console_handler = logging.StreamHandler()
    console_handler.setLevel(console_log_level)
    console_handler.setFormatter(log_format)
    handlers.append(console_handler)

    if log_file is not None and log_file != "":
        if os.path.exists(log_file):
            os.remove(log_file)
        elif not os.path.exists(os.path.dirname(log_file)):
            os.makedirs(os.path.dirname(log_file))
        file_handler = logging.FileHandler(log_file)
        file_handler.setLevel(log_file_level)
        file_handler.setFormatter(log_format)
        handlers.append(file_handler)

    logging.basicConfig(level=root_log_level, handlers=handlers)


class StoreLoggingLevelAction(configargparse.Action):
    """This class converts string into logging level"""

    LEVELS = {
        "CRITICAL": logging.CRITICAL,
        "ERROR": logging.ERROR,
        "WARNING": logging.WARNING,
        "INFO": logging.INFO,
        "DEBUG": logging.DEBUG,
        "NOTSET": logging.NOTSET,
    }

    CHOICES = list(LEVELS.keys()) + [str(_) for _ in LEVELS.values()]

    def __init__(self, option_strings, dest, help=None, **kwargs):
        super().__init__(option_strings, dest, help=help, **kwargs)

    def __call__(self, parser, namespace, value, option_string=None):
        """This function gets the key 'value' in the LEVELS, or just uses value"""

        level = StoreLoggingLevelAction.LEVELS.get(value, value)
        setattr(namespace, self.dest, level)


class CheckPathAction(configargparse.Action):
    """This class checks file path, if not exits, then create dir(file)"""

    def __init__(self, option_strings, dest, help=None, **kwargs):
        super().__init__(option_strings, dest, help=help, **kwargs)

    def __call__(self, parser, namespace, value, option_string=None):
        """This function checks file path, if not exits, then create dir(file)"""

        parent_path = os.path.dirname(value)
        if not os.path.exists(parent_path):
            os.makedirs(parent_path)
        setattr(namespace, self.dest, value)


class ConfigurationParer:
    """This class defines customized configuration parser"""

    def __init__(
        self,
        config_file_parser_class=configargparse.YAMLConfigFileParser,
        formatter_class=configargparse.ArgumentDefaultsHelpFormatter,
        **kwargs,
    ):
        """This funtion decides config parser and formatter

        Keyword Arguments:
            config_file_parser_class {configargparse.ConfigFileParser} -- config file parser (default: {configargparse.YAMLConfigFileParser})
            formatter_class {configargparse.ArgumentDefaultsHelpFormatter} -- config formatter (default: {configargparse.ArgumentDefaultsHelpFormatter})
        """

        self.parser = configargparse.ArgumentParser(
            config_file_parser_class=config_file_parser_class,
            formatter_class=formatter_class,
            **kwargs,
        )

    def add_save_cfgs(self):
        """This function adds saving path arguments: config file, model file..."""

        # config file configurations
        group = self.parser.add_argument_group("Config-File")
        group.add(
            "-config_file",
            "--config_file",
            required=False,
            is_config_file_arg=True,
            help="config file path",
        )

        # model file configurations
        group = self.parser.add_argument_group("Model-File")
        group.add(
            "-save_dir",
            "--save_dir",
            type=str,
            required=True,
            help="directory for saving checkpoints.",
        )

    def add_data_cfgs(self):
        """This function adds dataset arguments: data file path..."""

        self.parser.add(
            "-data_dir",
            "--data_dir",
            type=str,
            required=True,
            help="dataset directory.",
        )
        self.parser.add(
            "-train_file",
            "--train_file",
            type=str,
            required=False,
            help="train data file.",
        )
        self.parser.add(
            "-dev_file", "--dev_file", type=str, required=False, help="dev data file."
        )
        self.parser.add(
            "-test_file",
            "--test_file",
            type=str,
            required=False,
            help="test data file.",
        )
        self.parser.add(
            "-ent_rel_file",
            "--ent_rel_file",
            type=str,
            required=False,
            help="entity and relation file.",
        )
        self.parser.add(
            "-max_sent_len",
            "--max_sent_len",
            type=int,
            default=200,
            help="max sentence length.",
        )
        self.parser.add(
            "-max_wordpiece_len",
            "--max_wordpiece_len",
            type=int,
            default=512,
            help="max sentence length.",
        )
        self.parser.add("-test", "--test", action="store_true", help="testing mode")

    def add_model_cfgs(self):
        """This function adds model (network) arguments: embedding, hidden unit..."""

        # embedding configurations
        group = self.parser.add_argument_group("Embedding")
        group.add(
            "--freeze_bert",
            "--freeze_bert",
            action="store_true",
            help="freeze_bert",
        )
        group.add(
            "-load_weight_path",
            "--load_weight_path",
            type=str,
            default="",
            help="load_weight_path",
        )
        group.add(
            "-prune_topk",
            "--prune_topk",
            type=int,
            default=0,
            help="prune_topk",
        )
        group.add(
            "-task",
            "--task",
            type=str,
            choices=["quintuplet", "tagger", "triplet"],
            default="quintuplet",
            help="training task.",
        )
        group.add(
            "-embedding_model",
            "--embedding_model",
            type=str,
            choices=["bert", "pretrained"],
            default="bert",
            help="embedding model.",
        )
        group.add(
            "-bert_model_name",
            "--bert_model_name",
            type=str,
            required=False,
            help="bert model name.",
        )
        group.add(
            "-pretrained_model_name",
            "--pretrained_model_name",
            type=str,
            required=False,
            help="pretrained model name.",
        )
        group.add(
            "-bert_output_size",
            "--bert_output_size",
            type=int,
            default=768,
            help="bert output size.",
        )
        group.add(
            "-bert_dropout",
            "--bert_dropout",
            type=float,
            default=0.1,
            help="bert dropout rate.",
        )
        group.add(
            "--fine_tune",
            "--fine_tune",
            action="store_true",
            help="fine-tune pretrained model.",
        )

        # biaffine model
        group = self.parser.add_argument_group("Biaffine")
        group.add(
            "-max_span_length",
            "--max_span_length",
            type=int,
            default=10,
            help="maximum span length.",
        )
        group.add(
            "-mlp_hidden_size",
            "--mlp_hidden_size",
            type=int,
            default=768,
            help="mlp hidden units size.",
        )
        group.add(
            "-dropout", "--dropout", type=float, default=0.5, help="dropout rate."
        )
        group.add(
            "-separate_threshold",
            "--separate_threshold",
            type=float,
            default=1.4,
            help="the threshold for separating spans.",
        )
        group.add(
            "-logit_dropout",
            "--logit_dropout",
            type=float,
            default=0.1,
            help="logit dropout rate for robustness.",
        )

    def add_optimizer_cfgs(self):
        """This function adds optimizer arguments"""

        # gradient strategy
        self.parser.add(
            "-gradient_clipping",
            "--gradient_clipping",
            type=float,
            default=1.0,
            help="gradient clipping threshold.",
        )

        # learning rate
        self.parser.add(
            "--learning_rate",
            "-learning_rate",
            type=float,
            default=3e-5,
            help="Starting learning rate. "
            "Recommended settings: sgd = 1, adagrad = 0.1, "
            "adadelta = 1, adam = 0.001",
        )
        self.parser.add(
            "--bert_learning_rate",
            "-bert_learning_rate",
            type=float,
            default=3e-5,
            help="learning rate for bert, should be smaller than followed parts.",
        )
        self.parser.add(
            "-lr_decay_rate",
            "--lr_decay_rate",
            type=float,
            default=0.9,
            help="learn rate of layers decay rate.",
        )

        # Adam configurations
        group = self.parser.add_argument_group("Adam")
        group.add(
            "-adam_beta1",
            "--adam_beta1",
            type=float,
            default=0.9,
            help="The beta1 parameter used by Adam. "
            "Almost without exception a value of 0.9 is used in "
            "the literature, seemingly giving good results, "
            "so we would discourage changing this value from "
            "the default without due consideration.",
        )
        group.add(
            "-adam_beta2",
            "--adam_beta2",
            type=float,
            default=0.999,
            help="The beta2 parameter used by Adam. "
            "Typically a value of 0.999 is recommended, as this is "
            "the value suggested by the original paper describing "
            "Adam, and is also the value adopted in other frameworks "
            "such as Tensorflow and Kerras, i.e. see: "
            "https://www.tensorflow.org/api_docs/python/tf/train/Adam"
            "Optimizer or "
            "https://keras.io/optimizers/ . "
            'Whereas recently the paper "Attention is All You Need" '
            "suggested a value of 0.98 for beta2, this parameter may "
            "not work well for normal models / default "
            "baselines.",
        )
        group.add(
            "-adam_epsilon",
            "--adam_epsilon",
            type=float,
            default=1e-6,
            help="adam epsilon",
        )
        group.add(
            "-adam_weight_decay_rate",
            "--adam_weight_decay_rate",
            type=float,
            default=0.0,
            help="adam weight decay rate.",
        )
        group.add(
            "-adam_bert_weight_decay_rate",
            "--adam_bert_weight_decay_rate",
            type=float,
            default=0.0,
            help="adam weight decay rate of Bert module.",
        )

    def add_run_cfgs(self):
        """This function adds running arguments"""

        # training configurations
        group = self.parser.add_argument_group("Training")
        group.add("-seed", "--seed", type=int, default=5216, help="radom seed.")
        group.add(
            "-epochs", "--epochs", type=int, default=1000, help="training epochs."
        )
        group.add(
            "-pretrain_epochs",
            "--pretrain_epochs",
            type=int,
            default=0,
            help="pretrain epochs.",
        )
        group.add(
            "-warmup_rate",
            "--warmup_rate",
            type=float,
            default=0.0,
            help="warmup rate.",
        )
        group.add(
            "-early_stop",
            "--early_stop",
            type=int,
            default=50,
            help="early stop threshold.",
        )
        group.add(
            "-train_batch_size",
            "--train_batch_size",
            type=int,
            default=200,
            help="batch size during training.",
        )
        group.add(
            "-gradient_accumulation_steps",
            "--gradient_accumulation_steps",
            type=int,
            default=1,
            help="Number of updates steps to accumulate before performing a backward/update pass.",
        )

        # testing configurations
        group = self.parser.add_argument_group("Testing")
        group.add(
            "-test_batch_size",
            "--test_batch_size",
            type=int,
            default=100,
            help="batch size during testing.",
        )
        group.add(
            "-validate_every",
            "--validate_every",
            type=int,
            default=20000,
            help="output result every n samples during validating.",
        )

        # gpu configurations
        group = self.parser.add_argument_group("GPU")
        group.add(
            "-device",
            "--device",
            type=int,
            default=-1,
            help="cpu: device = -1, gpu: gpu device id(device >= 0).",
        )

        # logging configurations
        group = self.parser.add_argument_group("logging")
        group.add(
            "-root_log_level",
            "--root_log_level",
            type=str,
            action=StoreLoggingLevelAction,
            choices=StoreLoggingLevelAction.CHOICES,
            default="DEBUG",
            help="root logging out level.",
        )
        group.add(
            "-console_log_level",
            "--console_log_level",
            type=str,
            action=StoreLoggingLevelAction,
            choices=StoreLoggingLevelAction.CHOICES,
            default="NOTSET",
            help="console logging output level.",
        )
        group.add(
            "-log_file",
            "--log_file",
            type=str,
            required=True,
            help="logging file during running.",
        )
        group.add(
            "-file_log_level",
            "--file_log_level",
            type=str,
            action=StoreLoggingLevelAction,
            choices=StoreLoggingLevelAction.CHOICES,
            default="NOTSET",
            help="file logging output level.",
        )
        group.add(
            "-logging_steps",
            "--logging_steps",
            type=int,
            default=10,
            help="Logging every N update steps.",
        )

    def parse_args(self):
        """This function parses arguments and initializes logger

        Returns:
            dict -- config arguments
        """

        cfg = self.parser.parse_args()

        if not os.path.exists(cfg.save_dir):
            os.makedirs(cfg.save_dir)

        cfg.best_model_path = os.path.join(cfg.save_dir, "best_model")
        cfg.last_model_path = os.path.join(cfg.save_dir, "last_model")
        cfg.vocabulary_file = os.path.join(cfg.save_dir, "vocabulary.pickle")
        cfg.model_checkpoints_dir = os.path.join(cfg.save_dir, "model_ckpts")

        if not os.path.exists(cfg.model_checkpoints_dir):
            os.makedirs(cfg.model_checkpoints_dir)

        assert os.path.exists(
            cfg.data_dir
        ), f"dataset directory {cfg.data_dir} not exists !!!"
        for file in ["train_file", "dev_file", "test_file", "ent_rel_file"]:
            if getattr(cfg, file, None) is not None:
                setattr(cfg, file, os.path.join(cfg.data_dir, getattr(cfg, file, "")))

        if getattr(cfg, "log_file", None) is not None:
            cfg.log_file = os.path.join(cfg.save_dir, cfg.log_file)
            # assert not os.path.exists(cfg.log_file), f"log file {cfg.log_file} exists !!!"
            if os.path.exists(cfg.log_file):
                cfg.log_file = None
                cfg.log_file_level = logging.NOTSET

        init_logger(
            root_log_level=getattr(cfg, "root_log_level", logging.DEBUG),
            console_log_level=getattr(cfg, "console_log_level", logging.NOTSET),
            log_file=getattr(cfg, "log_file", None),
            log_file_level=getattr(cfg, "log_file_level", logging.NOTSET),
        )

        if cfg.freeze_bert:
            cfg.fine_tune = False
        return cfg

    def format_values(self):
        return self.parser.format_values()