from transformers import PretrainedConfig


class FLAMEConfig(PretrainedConfig):
    model_type = "flame"

    def __init__(
            self,
            patch_len: int = 48,
            expand: int = 2,
            d_conv: int = 4,
            d_model: int = 256,
            d_ff: int = 512,
            d_couple: int = 256,
            enc_layers: int = 1,
            dec_layers: int = 3,
            couple_layers: int = 5,
            head_dim: int = 64,
            n_heads: int = 4,
            activation: str = "gelu",
            dropout: float = 0.2,
            head_dropout: float = 0.2,
            learnable: bool = False,
            norm_mode: str = 'layer',
            **kwargs,
    ):
        self.patch_len = patch_len
        self.expand = expand
        self.d_conv = d_conv
        self.d_model = d_model
        self.d_ff = d_ff
        self.d_couple = d_couple
        self.enc_layers = enc_layers
        self.dec_layers = dec_layers
        self.couple_layers = couple_layers
        self.head_dim = head_dim
        self.n_heads = n_heads
        self.activation = activation
        self.dropout = dropout
        self.head_dropout = head_dropout
        self.learnable = learnable
        self.norm_mode = norm_mode

        super().__init__(
            **kwargs,
        )