Update modeling_spin.py

modeling_spin.py

@@ -108,8 +108,8 @@ class TransformerBlock(nn.Module):
 
 class SpinForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = SpinConfig
-    # Fix: Use dict mapping for tied weights
     _tied_weights_keys = {"lm_head.weight": "tok_embeddings.weight"}
+    _supports_cache_class = False
 
     def __init__(self, config: SpinConfig):
         super().__init__(config)
@@ -144,6 +144,7 @@ class SpinForCausalLM(PreTrainedModel, GenerationMixin):
     def forward(
         self,
         input_ids: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
         labels: torch.Tensor = None,
         past_key_values: list = None,
         use_cache: bool = False,
@@ -182,12 +183,13 @@ class SpinForCausalLM(PreTrainedModel, GenerationMixin):
             past_key_values=new_kv_caches if use_cache else None,
         )
 
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
         if past_key_values is not None and len(past_key_values) > 0 and past_key_values[0] is not None:
             input_ids = input_ids[:, -1:]
         return {
             "input_ids": input_ids,
             "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
             "use_cache": True,
         }