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=2048, 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
        self.num_hidden_layers = 1