from transformers import PretrainedConfig


class FelaTabConfig(PretrainedConfig):
    model_type = "felatab"

    def __init__(
        self,
        max_features=100,
        max_classes=10,
        dim=1024,
        n_layers=28,
        n_heads=16,
        head_dim=64,
        chunk=64,
        ffn_mult=3,
        dropout=0.0,
        use_landmark=True,
        n_landmarks=48,
        ln_eps=1e-05,
        **kwargs,
    ):
        self.max_features = max_features
        self.max_classes = max_classes
        self.dim = dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.chunk = chunk
        self.ffn_mult = ffn_mult
        self.dropout = dropout
        self.use_landmark = use_landmark
        self.n_landmarks = n_landmarks
        self.ln_eps = ln_eps
        super().__init__(**kwargs)