from transformers import PretrainedConfig

class CustomConfig(PretrainedConfig):
    
    model_type = "custom"

    def __init__(self,
                 max_sequence_length=None,
                 num_hidden_layers=None,
                 hidden_size=None,
                 use_parallel_embedding=False,
                 vocab_size=None,
                 position_embedding_type="ROPE",
                 is_neox_style=True,
                 num_attention_heads=None,
                 num_key_value_heads=None,
                 attention_bias=False,
                 intermediate_size=None,
                 hidden_act="silu",
                 is_gated=False,
                 num_experts=None,
                 num_experts_per_tok=None,
                 moe_intermediate_size=None,
                 shared_expert_intermediate_size=None,
                 norm_topk_prob=None,
                 mlp_bias=False,
                 norm_type="rmsnorm",
                 norm_eps=1e-05, 
                 apply_residual_connection_post_layernorm=False,
                 use_parallel_residual=False,
                 **kwargs):
        self.max_sequence_length = max_sequence_length
        self.num_hidden_layers = num_hidden_layers
        self.hidden_size = hidden_size
        self.use_parallel_embedding = use_parallel_embedding
        self.vocab_size = vocab_size
        self.position_embedding_type = position_embedding_type  # ALIBI, ROPE
        self.is_neox_style = is_neox_style  # True: fold_rotary; False: cross_rotary
        self.num_attention_heads = num_attention_heads
        if num_key_value_heads is None:
            self.num_key_value_heads = num_attention_heads
        else:
            self.num_key_value_heads = num_key_value_heads
        self.attention_bias = attention_bias
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act  # silu, gelu
        self.is_gated = is_gated
        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.moe_intermediate_size = moe_intermediate_size
        self.shared_expert_intermediate_size = shared_expert_intermediate_size
        self.norm_topk_prob = norm_topk_prob
        self.mlp_bias = mlp_bias
        self.norm_type = norm_type  # rmsnorm， layernorm
        self.norm_eps = norm_eps
        self.apply_residual_connection_post_layernorm = (
            apply_residual_connection_post_layernorm)
        self.use_parallel_residual = use_parallel_residual
        
        super().__init__(**kwargs)