from transformers import PretrainedConfig


class MyModelConfig(PretrainedConfig):
    model_type = "my_model"

    def __init__(
        self,
        n_layers=2,
        hidden_dim=3,
        **kwargs,
    ):
        self.n_layers = n_layers
        self.hidden_dim = hidden_dim
        super().__init__(**kwargs)