Create modelling_trm.py

modelling_trm.py

@@ -0,0 +1,123 @@
+import math
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from einops.layers.torch import EinMix
+from transformers import PreTrainedModel, PretrainedConfig
+
+# ---------------------------
+# Configuration Class
+# ---------------------------
+class TRMConfig(PretrainedConfig):
+    model_type = "trm"
+
+    def __init__(self,
+                 vocab_size=32000,
+                 hidden_size=256,
+                 seq_len=128,
+                 depth_L=2,
+                 depth_H=2,
+                 act_threshold=0.9,
+                 act_epsilon=1e-2,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.seq_len = seq_len
+        self.depth_L = depth_L
+        self.depth_H = depth_H
+        self.act_threshold = act_threshold
+        self.act_epsilon = act_epsilon
+
+
+# ---------------------------
+# Model Architecture
+# ---------------------------
+class HaltingBlock(nn.Module):
+    def __init__(self, hidden_size, act_threshold, act_epsilon):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size)
+        self.act_proj = nn.Linear(hidden_size, 1)
+        self.act_threshold = act_threshold
+        self.act_epsilon = act_epsilon
+
+    def forward(self, x):
+        halting_probs = torch.sigmoid(self.act_proj(x))
+        remainders = torch.zeros_like(halting_probs)
+        n_updates = torch.zeros_like(halting_probs)
+        still_running = torch.ones_like(halting_probs, dtype=torch.bool)
+        accumulated_output = torch.zeros_like(x)
+        accumulated_prob = torch.zeros_like(halting_probs)
+
+        while still_running.any():
+            p = torch.where(still_running, halting_probs, torch.zeros_like(halting_probs))
+            new_accum = accumulated_prob + p
+
+            still_running = new_accum < self.act_threshold
+            remainder = torch.where(still_running, torch.zeros_like(halting_probs), 1 - accumulated_prob)
+
+            update_weights = torch.where(still_running, p, remainder)
+            accumulated_output += update_weights * torch.tanh(self.proj(x))
+            accumulated_prob += update_weights
+            n_updates += still_running.float()
+
+            if (1 - accumulated_prob).mean() < self.act_epsilon:
+                break
+
+        return accumulated_output, accumulated_prob.mean()
+
+
+class TRMLayer(nn.Module):
+    def __init__(self, hidden_size, depth_H, act_threshold, act_epsilon):
+        super().__init__()
+        self.blocks = nn.ModuleList([
+            HaltingBlock(hidden_size, act_threshold, act_epsilon) for _ in range(depth_H)
+        ])
+        self.norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, x):
+        for block in self.blocks:
+            x, _ = block(x)
+        return self.norm(x)
+
+
+class TRM(PreTrainedModel):
+    config_class = TRMConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.pos_emb = nn.Parameter(torch.zeros(1, config.seq_len, config.hidden_size))
+        self.layers = nn.ModuleList([
+            TRMLayer(config.hidden_size, config.depth_H, config.act_threshold, config.act_epsilon)
+            for _ in range(config.depth_L)
+        ])
+        self.norm = nn.LayerNorm(config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    def forward(self, input_ids, labels=None):
+        x = self.emb(input_ids) + self.pos_emb[:, :input_ids.size(1), :]
+        for layer in self.layers:
+            x = layer(x)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        return {"loss": loss, "logits": logits}
+
+
+# ---------------------------
+# Utility: Register to AutoClasses
+# ---------------------------
+from transformers import AutoConfig, AutoModel
+
+AutoConfig.register("trm", TRMConfig)
+AutoModel.register(TRMConfig, TRM)