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