File size: 1,800 Bytes
7b96962 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | 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
|