from transformers import PretrainedConfig


class RiNALMoConfig(PretrainedConfig):
    model_type = "rinalmo"

    auto_map = {
        "AutoConfig": "configuration_rinalmo.RiNALMoConfig",
        "AutoModel": "modeling_rinalmo.RiNALMoModel",
        "AutoModelForMaskedLM": "modeling_rinalmo.RiNALMoForMaskedLM",
    }

    def __init__(
        self,
        vocab_size: int = 22,
        embed_dim: int = 1280,
        num_layers: int = 33,
        num_heads: int = 20,
        transition_factor: int = 4,
        padding_idx: int = 1,
        mask_idx: int = 4,
        cls_idx: int = 0,
        eos_idx: int = 2,
        unk_idx: int = 3,
        use_rot_emb: bool = True,
        rope_base: int = 10000,
        attention_dropout: float = 0.1,
        transition_dropout: float = 0.0,
        residual_dropout: float = 0.1,
        token_dropout_active: bool = True,
        mask_ratio: float = 0.15,
        mask_tkn_prob: float = 0.8,
        model_max_length: int = 8192,
        **kwargs,
    ):
        super().__init__(padding_idx=padding_idx, **kwargs)
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.transition_factor = transition_factor
        self.mask_idx = mask_idx
        self.cls_idx = cls_idx
        self.eos_idx = eos_idx
        self.unk_idx = unk_idx
        self.use_rot_emb = use_rot_emb
        self.rope_base = rope_base
        self.attention_dropout = attention_dropout
        self.transition_dropout = transition_dropout
        self.residual_dropout = residual_dropout
        self.token_dropout_active = token_dropout_active
        self.mask_ratio = mask_ratio
        self.mask_tkn_prob = mask_tkn_prob
        self.model_max_length = model_max_length