from .configuration_aliceai_t5 import AliceAIT5Config


class AliceAIT5MoEConfig(AliceAIT5Config):
    model_type = "aliceai_t5_moe"

    def __init__(
        self,
        expert_intermediate_size: int = 768,
        n_experts: int = 256,
        activation_type: str = "silu",
        routed_scaling_factor: float = 1.0,
        experts_n_group: int = 8,
        experts_top_k_group: int = 4,
        moe_normalize_expert_weights: int | float | bool | None = True,
        experts_gating_top: int = 4,
        group_routing: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.expert_intermediate_size = expert_intermediate_size
        self.n_experts = n_experts
        self.activation_type = activation_type
        self.routed_scaling_factor = routed_scaling_factor
        self.experts_n_group = experts_n_group
        self.experts_top_k_group = experts_top_k_group
        self.moe_normalize_expert_weights = moe_normalize_expert_weights
        self.experts_gating_top = experts_gating_top
        self.group_routing = group_routing