File size: 1,308 Bytes
e11caaf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | 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,
)
|