Update modeling_suprabrain.py

modeling_suprabrain.py

@@ -22,6 +22,7 @@ from torch.utils.data import Dataset, SequentialSampler
 
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers import GenerationMixin
 
 # ---------------------------------------------------------------------------
 # Optional backends
@@ -532,7 +533,7 @@ def _ce_chunk(h, W, labels, z_coef, softcap):
     return loss
 
 
-class SupraBrainForCausalLM(SupraBrainPreTrainedModel):
+class SupraBrainForCausalLM(SupraBrainPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {}
 
     def __init__(self, cfg: SupraBrainConfig):
@@ -543,6 +544,9 @@ class SupraBrainForCausalLM(SupraBrainPreTrainedModel):
             self.ub_a = nn.Parameter(torch.zeros(cfg.hidden_size, r))
             self.ub_b = nn.Parameter(torch.zeros(cfg.hidden_size, r))
             nn.init.normal_(self.ub_a, 0.0, cfg.initializer_range)
+        
+        # GenerationMixin requires main_input_name to be specified
+        self.main_input_name = "input_ids"
         self.post_init()
 
     def get_input_embeddings(self):
@@ -585,6 +589,16 @@ class SupraBrainForCausalLM(SupraBrainPreTrainedModel):
                 use_reentrant=False)
         return CausalLMOutputWithPast(loss=tot / ls.numel(), logits=None)
 
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, **kwargs
+    ):
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache", False),
+            "attention_mask": attention_mask,
+        }
+
 
 # ---- HF Auto-Registration -------------------------------------------------
 def register_hf():