from transformers import PretrainedConfig


class FelaModeratorConfig(PretrainedConfig):
    model_type = "fela-moderation"

    def __init__(
        self,
        vocab_size=259,
        max_len=512,
        d_model=448,
        n_layers=8,
        n_heads=7,
        fno_modes=128,
        gla_chunk=32,
        ffn_hidden=1280,
        layer_pattern="SSSL",
        dropout=0.1,
        pad_id=256,
        n_tox_labels=6,
        n_pii_tags=113,
        n_tax=19,
        n_spam=3,
        n_jailbreak=4,
        n_nsfw=2,
        n_identity=7,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.max_len = max_len
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.fno_modes = fno_modes
        self.gla_chunk = gla_chunk
        self.ffn_hidden = ffn_hidden
        self.layer_pattern = layer_pattern
        self.dropout = dropout
        self.pad_id = pad_id
        self.n_tox_labels = n_tox_labels
        self.n_pii_tags = n_pii_tags
        self.n_tax = n_tax
        self.n_spam = n_spam
        self.n_jailbreak = n_jailbreak
        self.n_nsfw = n_nsfw
        self.n_identity = n_identity
        super().__init__(**kwargs)