from transformers import PretrainedConfig


class RNAMSMConfig(PretrainedConfig):
    model_type = "rnamsm"

    auto_map = {
        "AutoConfig": "configuration_rnamsm.RNAMSMConfig",
        "AutoModel": "modeling_rnamsm.RNAMSMModel",
        "AutoModelForMaskedLM": "modeling_rnamsm.RNAMSMForMaskedLM",
    }

    def __init__(
        self,
        vocab_size=12,
        num_layers=10,
        embed_dim=768,
        num_attention_heads=12,
        ffn_embed_dim=3072,
        padding_idx=1,
        mask_idx=11,
        cls_idx=0,
        eos_idx=2,
        dropout=0.1,
        attention_dropout=0.1,
        activation_dropout=0.1,
        max_positions=1024,
        max_alignments=1024,
        max_tokens_per_msa=16384,
        embed_positions_msa=True,
        **kwargs,
    ):
        super().__init__(padding_idx=padding_idx, **kwargs)
        self.vocab_size = vocab_size
        self.num_layers = num_layers
        self.embed_dim = embed_dim
        self.num_attention_heads = num_attention_heads
        self.ffn_embed_dim = ffn_embed_dim
        self.mask_idx = mask_idx
        self.cls_idx = cls_idx
        self.eos_idx = eos_idx
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.activation_dropout = activation_dropout
        self.max_positions = max_positions
        self.max_alignments = max_alignments
        self.max_tokens_per_msa = max_tokens_per_msa
        self.embed_positions_msa = embed_positions_msa