from transformers import PretrainedConfig


class RabbitConfig(PretrainedConfig):
    model_type = "rabbit_ssm"

    def __init__(
        self,
        vocab_size: int = 50280,
        d_model: int = 2560,
        n_layers: int = 64,
        **kwargs,
    ):
        self.d_model = d_model
        self.n_layers = n_layers
        super().__init__(vocab_size=vocab_size, **kwargs)