File size: 828 Bytes
ae69998 2bb7fc4 ae69998 | 1 2 3 4 5 6 7 8 9 10 |
from transformers import PretrainedConfig
class TRMTextISMConfig(PretrainedConfig):
model_type = "trm_text_ism"
def __init__(self, vocab_size=50257, max_seq_len=512, dim=768, n_heads=12, head_dim=64, recurrence_steps=4, mlp_ratio=2.6666666667, mlp_hidden_size=None, dropout=0.0, gate_style="stable", gate_init=-1.5, residual_scale=0.5, tie_word_embeddings=True, **kwargs):
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
self.vocab_size, self.max_seq_len, self.dim, self.n_heads, self.head_dim = vocab_size, max_seq_len, dim, n_heads, head_dim
self.recurrence_steps, self.mlp_ratio, self.mlp_hidden_size, self.dropout = recurrence_steps, mlp_ratio, mlp_hidden_size, dropout
self.gate_style, self.gate_init, self.residual_scale = gate_style, gate_init, residual_scale
|