from transformers import PretrainedConfig

class SovythosConfig(PretrainedConfig):
    model_type = "sovythos"

    def __init__(
        self,
        vocab_size=32000,
        dim=1024,
        n_layers=24,
        n_heads=16,
        n_kv_heads=4,
        max_seq_len=2048,
        rope_theta=1000000.0,
        norm_eps=1e-06,
        ffn_multiple_of=256,
        ffn_dim_multiplier=None,
        tie_embeddings=True,
        dropout=0.0,
        use_grad_checkpoint=False,
        identity="SOVYTHOS-V2",
        bos_token_id=0,
        eos_token_id=0,
        pad_token_id=1,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.dim = dim
        self.hidden_size = dim
        self.n_layers = n_layers
        self.num_hidden_layers = n_layers
        self.n_heads = n_heads
        self.num_attention_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.num_key_value_heads = n_kv_heads
        self.max_seq_len = max_seq_len
        self.max_position_embeddings = max_seq_len
        self.rope_theta = rope_theta
        self.norm_eps = norm_eps
        self.rms_norm_eps = norm_eps
        self.ffn_multiple_of = ffn_multiple_of
        self.ffn_dim_multiplier = ffn_dim_multiplier
        self.tie_embeddings = tie_embeddings
        self.dropout = dropout
        self.use_grad_checkpoint = use_grad_checkpoint
        self.identity = identity
        kwargs.pop("tie_word_embeddings", None)
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_embeddings,
            **kwargs,
        )