"""TRM MoE configuration."""
from transformers import PretrainedConfig

class TRMMoEConfig(PretrainedConfig):
    model_type = "trm_moe"

    def __init__(
        self,
        vocab_size=50259,
        dim=768,
        n_heads=12,
        head_dim=64,
        mlp_ratio=2.6666666667,
        mlp_hidden_size=None,
        num_hidden_layers=1,
        recurrence_steps=4,
        residual_scale=0.5,
        gate_style="stable",
        gate_init=-1.5,
        dropout=0.0,
        max_seq_len=512,
        num_experts=32,
        num_experts_per_token=2,
        moe_aux_loss_coeff=0.01,
        moe_z_loss_coeff=0.001,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.mlp_ratio = mlp_ratio
        self.mlp_hidden_size = mlp_hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.recurrence_steps = recurrence_steps
        self.residual_scale = residual_scale
        self.gate_style = gate_style
        self.gate_init = gate_init
        self.dropout = dropout
        self.max_seq_len = max_seq_len
        self.num_experts = num_experts
        self.num_experts_per_token = num_experts_per_token
        self.moe_aux_loss_coeff = moe_aux_loss_coeff
        self.moe_z_loss_coeff = moe_z_loss_coeff