from transformers import PretrainedConfig


class TinyIndianLMConfig(PretrainedConfig):
    model_type = "tinyindianlm"

    def __init__(
        self,
        vocab_size=16000,
        d_model=512,
        n_layers=4,
        n_heads=8,
        n_kv_heads=4,
        max_seq_len=512,
        d_ffn=1408,
        dropout=0.0,
        rope_theta=10000.0,
        pad_token_id=0,
        unk_token_id=1,
        bos_token_id=2,
        eos_token_id=3,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.max_seq_len = max_seq_len
        self.d_ffn = d_ffn
        self.dropout = dropout
        self.rope_theta = rope_theta
        super().__init__(
            pad_token_id=pad_token_id,
            unk_token_id=unk_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )