Upload folder using huggingface_hub

chat_template.jinja

@@ -0,0 +1 @@
+{% for message in messages %}{{'<|' + message['role'] + '|>\n' + message['content'] + '<|end|>\n'}}{% endfor %}

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "TRMTextISMForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_trm_text_ism.TRMTextISMConfig",
+    "AutoModelForCausalLM": "modeling_trm_text_ism.TRMTextISMForCausalLM"
+  },
+  "bos_token_id": 50256,
+  "dim": 768,
+  "dropout": 0.0,
+  "dtype": "bfloat16",
+  "eos_token_id": 50256,
+  "gate_init": -1.5,
+  "gate_style": "stable",
+  "head_dim": 64,
+  "max_seq_len": 512,
+  "mlp_hidden_size": null,
+  "mlp_ratio": 2.6666666667,
+  "model_type": "trm_text_ism",
+  "n_heads": 12,
+  "num_hidden_layers": 1,
+  "pad_token_id": 50256,
+  "recurrence_steps": 4,
+  "residual_scale": 0.5,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.9.0",
+  "use_cache": false,
+  "vocab_size": 50257
+}

configuration_trm_text_ism.py

@@ -0,0 +1,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=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

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 50256,
+  "do_sample": true,
+  "eos_token_id": [
+    50256,
+    50256
+  ],
+  "max_new_tokens": 128,
+  "pad_token_id": 50256,
+  "temperature": 0.8,
+  "top_k": 50,
+  "top_p": 0.95,
+  "transformers_version": "5.9.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3654eea5d60bd4614874eabab92f00b295f90ac3b4561c5a08668d8f7aedafe2
+size 168691448

modeling_trm_text_ism.py

@@ -0,0 +1,82 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from .configuration_trm_text_ism import TRMTextISMConfig
+
+def apply_rope(x, cos, sin):
+    S = x.shape[2]
+    c, s = cos[:, :, :S, :].to(x.dtype), sin[:, :, :S, :].to(x.dtype)
+    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        h = config.mlp_hidden_size or int(config.dim * config.mlp_ratio)
+        self.gate_proj = nn.Linear(config.dim, h, bias=False)
+        self.up_proj = nn.Linear(config.dim, h, bias=False)
+        self.down_proj = nn.Linear(h, config.dim, bias=False)
+    def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+class TRMAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.n_heads, self.head_dim = config.n_heads, config.head_dim
+        self.qkv = nn.Linear(config.dim, 3*config.dim, bias=False)
+        self.out = nn.Linear(config.dim, config.dim, bias=False)
+    def forward(self, x, mask, cos, sin):
+        B, S, _ = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        q, k, v = [t.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) for t in (q, k, v)]
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None, :, :])
+        return self.out(y.transpose(1, 2).reshape(B, S, -1))
+
+class TRMBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.res = config.residual_scale
+        self.norm1 = nn.RMSNorm(config.dim)
+        self.attn = TRMAttention(config)
+        self.norm2 = nn.RMSNorm(config.dim)
+        self.mlp = SwiGLUMLP(config)
+        self.attn_gate = nn.Parameter(torch.ones(config.dim))
+        self.mlp_gate = nn.Parameter(torch.ones(config.dim))
+    def forward(self, x, mask, c, s):
+        x = x + self.res * torch.sigmoid(self.attn_gate).view(1,1,-1) * self.attn(self.norm1(x), mask, c, s)
+        return x + self.res * torch.sigmoid(self.mlp_gate).view(1,1,-1) * self.mlp(self.norm2(x))
+
+class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = TRMTextISMConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.token_emb = nn.Embedding(config.vocab_size, config.dim)
+        self.block = TRMBlock(config)
+        self.norm = nn.RMSNorm(config.dim)
+        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
+        pos = torch.arange(config.max_seq_len).float()
+        theta = 1.0 / (10000.0 ** (torch.arange(0, config.head_dim//2).float() / (config.head_dim//2)))
+        f = torch.outer(pos, theta)
+        self.register_buffer("rope_cos", f.cos().view(1, 1, config.max_seq_len, -1))
+        self.register_buffer("rope_sin", f.sin().view(1, 1, config.max_seq_len, -1))
+        self.post_init()
+    def get_input_embeddings(self): return self.token_emb
+    def get_output_embeddings(self): return self.lm_head
+    def tie_weights(self, *args, **kwargs):
+        if hasattr(self, 'lm_head'): self.lm_head.weight = self.token_emb.weight
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        B, S = input_ids.shape
+        x = self.token_emb(input_ids)
+        m = torch.tril(torch.ones(S, S, device=input_ids.device)).bool().unsqueeze(0).expand(B, -1, -1)
+        if attention_mask is not None: m = m & attention_mask[:, None, :].bool()
+        c, s = self.rope_cos, self.rope_sin
+        for _ in range(self.config.recurrence_steps):
+            x = self.block(x, m, c, s)
+        logits = self.lm_head(self.norm(x))
+        return CausalLMOutputWithPast(logits=logits)

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a890963eda01ff5a4e105510bc22992df7bd6714a4f50cbd20f9ab23fa7f65e
+size 5265