from transformers import PretrainedConfig


class SpinConfig(PretrainedConfig):
    model_type = "spin"
    attribute_map = {
        "hidden_size": "d_model",
        "num_hidden_layers": "n_layers",
        "num_attention_heads": "n_heads",
    }

    def __init__(
        self,
        vocab_size: int = 512,
        max_seq_len: int = 256,
        d_model: int = 48,
        n_layers: int = 2,
        n_heads: int = 4,
        d_ff: int = 128,
        norm_eps: float = 1e-5,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.d_ff = d_ff
        self.norm_eps = norm_eps

        
        self.hidden_size = d_model
        self.num_hidden_layers = n_layers
        self.num_attention_heads = n_heads

        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)