from transformers import PretrainedConfig

class PebbleConfig(PretrainedConfig):
    model_type = "pebble_25m"

    def __init__(
        self,
        vocab_size=2048,
        d_model=608,
        n_blocks=8,
        n_heads=8,
        att_hidden=2432,
        rope_theta=10000.0,
        seq_len=2048,
        mamba_d_state=128,
        mamba_d_conv=4,
        mamba_expand=2,
        mamba_headdim=64,
        block_pattern="3:1 mamba:attn",
        bos_token_id=1,
        eos_token_id=1,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_blocks = n_blocks
        self.n_heads = n_heads
        self.att_hidden = att_hidden
        self.rope_theta = rope_theta
        self.seq_len = seq_len
        self.mamba_d_state = mamba_d_state
        self.mamba_d_conv = mamba_d_conv
        self.mamba_expand = mamba_expand
        self.mamba_headdim = mamba_headdim
        self.block_pattern = block_pattern
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )