| 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 | |