from transformers import PretrainedConfig


class FelaConfig(PretrainedConfig):
    model_type = "fela"

    def __init__(
        self,
        vocab_size=151936,
        seq_len=2048,
        n_layer=28,
        n_embd=1536,
        n_head=12,
        ffn_hidden=8960,
        layer_pattern="SSSL",
        gla_delta=True,
        fno_modes=512,
        gla_chunk=256,
        landmark_layer_every=7,
        landmark_chunk=32,
        landmark_max=64,
        attn_layer_every=0,
        tie_word_embeddings=False,
        use_cache=False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.n_layer = n_layer
        self.n_embd = n_embd
        self.n_head = n_head
        self.ffn_hidden = ffn_hidden
        self.layer_pattern = layer_pattern
        self.gla_delta = gla_delta
        self.fno_modes = fno_modes
        self.gla_chunk = gla_chunk
        self.landmark_layer_every = landmark_layer_every
        self.landmark_chunk = landmark_chunk
        self.landmark_max = landmark_max
        self.attn_layer_every = attn_layer_every
        self.num_hidden_layers = n_layer
        self.use_cache = use_cache
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)