from transformers import PretrainedConfig


class FelaTsConfig(PretrainedConfig):
    model_type = "fela-ts"

    def __init__(
        self,
        C=321,
        L=512,
        H=96,
        patch=16,
        stride=8,
        D=128,
        modes=16,
        nblk=3,
        **kwargs,
    ):
        self.C = C
        self.L = L
        self.H = H
        self.patch = patch
        self.stride = stride
        self.D = D
        self.modes = modes
        self.nblk = nblk
        super().__init__(**kwargs)