from transformers import LlamaConfig

__all__ = ["HFOpenMoeConfig"]


class HFOpenMoeConfig(LlamaConfig):
    model_type = "openmoe"
    def __init__(
            self,
            num_experts: int = 32,
            moe_layer_interval: int = 6,
            router_topk: int = 2,
            router_capacity_factor_train: float = 1.25,
            router_capacity_factor_eval: float = 2.0,
            router_min_capacity: int = 4,
            router_noisy_policy: str = None,
            router_drop_tks: bool = True,
            router_aux_loss_factor: float = 0.01,
            router_z_loss_factor: float = 0.0001,
            mlp_gated: bool = True,
            label_smoothing: float = 0.001,
            z_loss_factor: float = 0.01,
            enable_load_balance: bool = False,
            load_balance_tolerance: float = 0.1,
            load_balance_beam_width: int = 8,
            load_balance_group_swap_factor: float = 0.4,
            enable_kernel: bool = False,
            enable_comm_overlap: bool = False,
            enable_hierarchical_alltoall: bool = False,
            **kwargs
    ):
        self.num_experts = num_experts
        self.moe_layer_interval = moe_layer_interval
        self.router_topk = router_topk
        self.router_capacity_factor_train = router_capacity_factor_train
        self.router_capacity_factor_eval = router_capacity_factor_eval
        self.router_min_capacity = router_min_capacity
        self.router_noisy_policy = router_noisy_policy
        self.router_drop_tks = router_drop_tks
        self.router_aux_loss_factor = router_aux_loss_factor
        self.router_z_loss_factor = router_z_loss_factor
        self.mlp_gated = mlp_gated
        self.label_smoothing = label_smoothing
        self.z_loss_factor = z_loss_factor
        self.enable_load_balance = enable_load_balance
        self.load_balance_tolerance = load_balance_tolerance
        self.load_balance_beam_width = load_balance_beam_width
        self.load_balance_group_swap_factor = load_balance_group_swap_factor
        self.enable_kernel = enable_kernel
        self.enable_comm_overlap = enable_comm_overlap
        self.enable_hierarchical_alltoall = enable_hierarchical_alltoall

        super().__init__(**kwargs)