| 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, | |
| ) | |