Peft model signature

src/peft/__init__.py

@@ -68,4 +68,5 @@
     prepare_model_for_kbit_training,
     set_peft_model_state_dict,
     shift_tokens_right,
+    update_forward_signature,
 )

src/peft/peft_model.py

@@ -77,6 +77,10 @@
 }
 
 
+def forward(self, *args, **kwargs):
+    return self.get_base_model()(*args, **kwargs)
+
+
 class PeftModel(PushToHubMixin, torch.nn.Module):
     """
     Base model encompassing various Peft methods.
@@ -109,6 +113,7 @@ def __init__(self, model: PreTrainedModel, peft_config: PeftConfig, adapter_name
         self.peft_config = {}
         self.active_adapter = adapter_name
         self.peft_type = peft_config.peft_type
+        self.parent_class = type(model)  # Gets the parent class of the original model to update signature
         if not isinstance(peft_config, PromptLearningConfig):
             self.peft_config[adapter_name] = peft_config
             self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type](

src/peft/utils/__init__.py

@@ -43,6 +43,7 @@
     ModulesToSaveWrapper,
     _prepare_prompt_learning_config,
     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

src/peft/utils/other.py

@@ -16,12 +16,15 @@
 import inspect
 import os
 import warnings
-from typing import Optional
+from functools import update_wrapper
+from types import MethodType
+from typing import Optional, Union
 
 import accelerate
 import torch
 from accelerate.hooks import add_hook_to_module, remove_hook_from_module
 from accelerate.utils import is_npu_available, is_xpu_available
+from transformers.utils import PushToHubMixin
 
 
 # Get current device name based on available devices
@@ -313,6 +316,28 @@ 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) -> Union[torch.nn.Module, PushToHubMixin]:
+    """
+    Updates the forward signature of the PeftModel to include parents class signature
+    """
+    # 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, model.parent_class.forward, assigned=("__doc__", "__name__", "__annotations__")
+        )
+        model.forward = MethodType(default_forward, model)
+    return model
+
+
 TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
     "t5": ["q", "v"],
     "mt5": ["q", "v"],