Update modeling_medassistgpt.py

modeling_medassistgpt.py

@@ -105,6 +105,20 @@ class MedAssistGPTModel(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, **kwargs
+    ):
+        # If past_key_values is provided, only use the last token
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": use_cache,
+            "attention_mask": attention_mask,
+        }
+
     def forward(self, input_ids, labels=None):
         h = self.embed(input_ids)
         for blk in self.blocks: