from transformers import PretrainedConfig


class RnaFmConfig(PretrainedConfig):
    model_type = "rnafm"

    auto_map = {
        "AutoConfig": "configuration_rnafm.RnaFmConfig",
        "AutoModel": "modeling_rnafm.RnaFmModel",
        "AutoModelForMaskedLM": "modeling_rnafm.RnaFmForMaskedLM",
        "AutoTokenizer": ["tokenization_rnafm.RnaFmTokenizer", None],
    }

    def __init__(
        self,
        vocab_size: int = 25,
        num_layers: int = 12,
        embed_dim: int = 640,
        ffn_embed_dim: int = 5120,
        attention_heads: int = 20,
        padding_idx: int = 1,
        mask_idx: int = 24,
        cls_idx: int = 0,
        eos_idx: int = 2,
        token_dropout: bool = False,
        emb_layer_norm_before: bool = True,
        model_max_length: int = 1024,
        model_variant: str = "rna",
        **kwargs,
    ):
        super().__init__(padding_idx=padding_idx, **kwargs)
        self.vocab_size = vocab_size
        self.num_layers = num_layers
        self.embed_dim = embed_dim
        self.ffn_embed_dim = ffn_embed_dim
        self.attention_heads = attention_heads
        self.mask_idx = mask_idx
        self.cls_idx = cls_idx
        self.eos_idx = eos_idx
        self.token_dropout = token_dropout
        self.emb_layer_norm_before = emb_layer_norm_before
        self.model_max_length = model_max_length
        self.model_variant = model_variant