Peft model signature

src/peft/__init__.py

@@ -72,6 +72,7 @@
     prepare_model_for_kbit_training,
     set_peft_model_state_dict,
     shift_tokens_right,
+    update_forward_signature,
     load_peft_weights,
 )
 from .config import PeftConfig, PromptLearningConfig

src/peft/utils/__init__.py

@@ -45,6 +45,7 @@
     _prepare_prompt_learning_config,
     _is_valid_match,
     infer_device,
+    update_forward_signature,
 )
 from .hub_utils import hub_file_exists
 from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict, load_peft_weights

src/peft/utils/other.py

@@ -16,6 +16,8 @@
 import inspect
 import os
 import warnings
+from functools import update_wrapper
+from types import MethodType
 from typing import Optional
 
 import accelerate
@@ -325,6 +327,42 @@ def _get_batch_size(input_ids: Optional[torch.Tensor], inputs_embeds: Optional[t
     return batch_size
 
 
+def default_forward(self, *args, **kwargs):
+    return self.get_base_model()(*args, **kwargs)
+
+
+def update_forward_signature(model):
+    """
+    Updates the forward signature of the PeftModel to include parents class signature
+    Args:
+        model (`PeftModel`): model
+    Example:
+
+    ```python
+    >>> from transformers import  WhisperForConditionalGeneration
+    >>> from peft import  get_peft_model, LoraConfig,  update_forward_signature
+
+    >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en")
+    >>> peft_config = LoraConfig(r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["q_proj", "v_proj"])
+
+    >>> peft_model = get_peft_model(model, peft_config)
+    >>> update_forward_signature(peft_model)
+    ```
+    """
+
+    # Only update signature when the current forward signature only has *args and **kwargs
+    current_signature = inspect.signature(model.forward)
+    if (
+        len(current_signature.parameters) == 2
+        and "args" in current_signature.parameters
+        and "kwargs" in current_signature.parameters
+    ):
+        update_wrapper(
+            default_forward, type(model.get_base_model()).forward, assigned=("__doc__", "__name__", "__annotations__")
+        )
+        model.forward = MethodType(default_forward, model)
+
+
 TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
     "t5": ["q", "v"],
     "mt5": ["q", "v"],