search_weight_scale_config.py Source File

search_weight_scale_config.py Source File#

Mobilint SDK qb Compiler: search_weight_scale_config.py Source File
Mobilint SDK qb Compiler v0.12.0.0
MCS002-KR
from typeguard import typechecked
from . import ConfigABC
 
 
class SearchWeightScaleConfig(ConfigABC):
    """
    @brief Configuration for searching weight scale.
    """
 
    optional = False
    DEFAULTS = {
        "apply": False,
        "query": False,
        "key": False,
        "value": False,
        "out": False,
        "ffn": False,
    }
 
    def __init__(
        self,
        apply: bool,
        query: bool,
        key: bool,
        value: bool,
        out: bool,
        ffn: bool,
        optional: bool = optional,
    ):
        """
        @brief Initialize the SearchWeightScaleConfig.
        @param apply bool. Whether to apply weight scale search.
        @param query bool. Apply to query layer.
        @param key bool. Apply to key layer.
        @param value bool. Apply to value layer.
        @param out bool. Apply to output layer.
        @param ffn bool. Apply to FFN layer.
        @param optional bool. Indicates whether this config is optional. Use the default value defined as a class variable.
        """
        super().__init__(optional)
        self.apply = apply
        self.query = query
        self.key = key
        self.value = value
        self.out = out
        self.ffn = ffn
 
    @classmethod
    def default_config(cls):
        return cls(**cls.DEFAULTS, optional=cls.optional)
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        """Create SearchWeightScaleConfig from dictionary (JSON structure)"""
        params = cls.DEFAULTS.copy()
 
        if "apply" in config_dict:
            params["apply"] = config_dict["apply"]
 
        if "transformer" in config_dict:
            transformer = config_dict["transformer"]
            if "query" in transformer:
                params["query"] = transformer["query"]
            if "key" in transformer:
                params["key"] = transformer["key"]
            if "value" in transformer:
                params["value"] = transformer["value"]
            if "out" in transformer:
                params["out"] = transformer["out"]
            if "ffn" in transformer:
                params["ffn"] = transformer["ffn"]
 
        return cls(**params)
 
    def to_dict(self):
        return {
            "apply": self.apply,
            "transformer": {
                "query": self.query,
                "key": self.key,
                "value": self.value,
                "out": self.out,
                "ffn": self.ffn,
            },
        }
 
 
def get_search_weight_scale_config(**kwargs) -> SearchWeightScaleConfig:
    """
    @brief Create SearchWeightScaleConfig with partial parameters merged with defaults.
    """
    params = SearchWeightScaleConfig.DEFAULTS.copy()
    params.update(kwargs)
    return SearchWeightScaleConfig(**params, optional=SearchWeightScaleConfig.optional)
 
 
# }@