From 5bc815e2e2e484ae89c103d9fc624f4e53e05327 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Mon, 6 Feb 2023 23:50:48 +0530
Subject: [PATCH] fix `generate` because of recent transformers release

---
 src/peft/peft_model.py | 25 ++++++++-----------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index ea24090..4dfec27 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -539,16 +539,16 @@ class PeftModelForCausalLM(PeftModel):
                 kwargs["token_type_ids"] = None
 
             if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
-                batch_size = kwargs["input_ids"].shape[0]
-                past_key_values = self.get_prompt(batch_size)
-                kwargs["past_key_values"] = past_key_values
                 return self.base_model.generate(**kwargs)
             else:
                 raise NotImplementedError
 
     def prepare_inputs_for_generation(self, *args, **kwargs):
         model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
-        model_kwargs["past_key_values"] = kwargs.get("past", None) or kwargs.get("past_key_values", None)
+        if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            batch_size = model_kwargs["input_ids"].shape[0]
+            past_key_values = self.get_prompt(batch_size)
+            model_kwargs["past_key_values"] = past_key_values
         return model_kwargs
 
 
@@ -682,25 +682,16 @@ class PeftModelForSeq2SeqLM(PeftModel):
                 kwargs["token_type_ids"] = None
 
             if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
-                batch_size = kwargs["input_ids"].shape[0]
-                past_key_values = self.get_prompt(batch_size)
-                kwargs["past_key_values"] = past_key_values
                 return self.base_model.generate(**kwargs)
             else:
                 raise NotImplementedError
 
     def prepare_inputs_for_generation(self, *args, **kwargs):
         model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
-        model_kwargs["past_key_values"] = kwargs.get("past", None) or kwargs.get("past_key_values", None)
-        return model_kwargs
-
-    def _prepare_encoder_decoder_kwargs_for_generation(self, inputs_tensor, model_kwargs, model_input_name=None):
-        past_key_values = model_kwargs.get("past_key_values", None)
-        model_kwargs["past_key_values"] = None
-        model_kwargs = self.base_model_prepare_encoder_decoder_kwargs_for_generation(
-            inputs_tensor, model_kwargs, model_input_name
-        )
-        model_kwargs["past_key_values"] = past_key_values
+        if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            batch_size = model_kwargs["decoder_input_ids"].shape[0]
+            past_key_values = self.get_prompt(batch_size)
+            model_kwargs["past_key_values"] = past_key_values
         return model_kwargs