from transformers import PretrainedConfig


class RNETinyGPTConfig(PretrainedConfig):
    model_type = "rne_tiny_gpt"

    def __init__(
        self,
        vocab_size=32768,
        ctx_len=4096,
        n_layer=4,
        n_head=4,
        n_embd=384,
        dropout=0.0,
        pad_token_id=0,
        sep_token_id=3,
        pooling="mean",
        normalize_embeddings=True,
        attention_backend="sage",
        torch_fallback=False,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            sep_token_id=sep_token_id,
            **kwargs,
        )

        self.vocab_size = int(vocab_size)
        self.ctx_len = int(ctx_len)
        self.max_position_embeddings = int(ctx_len)

        self.n_layer = int(n_layer)
        self.n_head = int(n_head)
        self.n_embd = int(n_embd)

        self.num_hidden_layers = int(n_layer)
        self.num_attention_heads = int(n_head)
        self.hidden_size = int(n_embd)

        self.dropout = float(dropout)
        self.pooling = str(pooling)
        self.normalize_embeddings = bool(normalize_embeddings)
        self.attention_backend = str(attention_backend)
        self.torch_fallback = bool(torch_fallback)