diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 14dd6f2..68de0c2 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -21,7 +21,7 @@ from .peft_model import ( PeftModelForTokenClassification, ) from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig -from .utils import PeftType, PromptLearningConfig +from .utils import PromptLearningConfig MODEL_TYPE_TO_PEFT_MODEL_MAPPING = { diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 7403d6b..ef8892a 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -245,13 +245,20 @@ class PeftModel(PushToHubMixin, torch.nn.Module): def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" - if name == "forward": - return getattr(self.base_model, name) try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: return getattr(self.base_model, name) + def forward(self, *args, **kwargs): + """ + Forward pass of the model. + """ + if isinstance(self.peft_config, PromptLearningConfig): + return self.base_model(*args, **kwargs) + else: + return self.base_model.model(*args, **kwargs) + class PeftModelForSequenceClassification(PeftModel): """ diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index f4676d2..567ecc6 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -115,6 +115,7 @@ class LoraModel(torch.nn.Module): self.model = model self._find_and_replace() mark_only_lora_as_trainable(self.model, self.peft_config.bias) + self.forward = self.model.forward def _find_and_replace(self): kwargs = { @@ -173,8 +174,6 @@ class LoraModel(torch.nn.Module): def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" - if name == "forward": - return getattr(self.model, name) try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: