Upload modeling_trm_text_ism.py

modeling_trm_text_ism.py

@@ -1,4 +1,3 @@
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -64,17 +63,32 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
         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 get_input_embeddings(self):
+        return self.token_emb
+
+    def set_input_embeddings(self, value):
+        self.token_emb = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
     def tie_weights(self, *args, **kwargs):
-        if hasattr(self, 'lm_head'): self.lm_head.weight = self.token_emb.weight
+        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()
+        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)