from transformers import PretrainedConfig


class axiomConfig(PretrainedConfig):
    model_type = "axiom"

    def __init__(
        self,
        vocab_size=100277,
        dim=1024,
        n_layers=24,
        n_heads=16,
        n_kv_heads=8,
        ffn_dim_multiplier=2.6667,
        max_seq_len=1024,
        rope_theta=10000.0,
        norm_eps=1e-5,
        dropout=0.0,
        bos_token_id=None,
        eos_token_id=100257,
        pad_token_id=100257,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.ffn_dim_multiplier = ffn_dim_multiplier
        self.max_seq_len = max_seq_len
        self.rope_theta = rope_theta
        self.norm_eps = norm_eps
        self.dropout = dropout
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            **kwargs,
        )