diff --git a/src/peft/mapping.py b/src/peft/mapping.py index e885ea4..14dd6f2 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 PromptLearningConfig +from .utils import PeftType, PromptLearningConfig MODEL_TYPE_TO_PEFT_MODEL_MAPPING = { @@ -135,8 +135,11 @@ def get_peft_model(model, peft_config): model_config = model.config.to_dict() peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) - if not isinstance(peft_config, PromptLearningConfig): + if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): peft_config = _prepare_lora_config(peft_config, model_config) return PeftModel(model, peft_config) - peft_config = _prepare_prompt_learning_config(peft_config, model_config) + if not isinstance(peft_config, PromptLearningConfig): + peft_config = _prepare_lora_config(peft_config, model_config) + else: + peft_config = _prepare_prompt_learning_config(peft_config, model_config) return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 38993bf..7403d6b 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -76,7 +76,6 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if getattr(self.peft_config, "modules_to_save", None) is not None: self.modules_to_save = self.peft_config.modules_to_save _set_trainable(self) - self.forward = self.base_model.forward self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def save_pretrained(self, save_directory, **kwargs): @@ -246,6 +245,8 @@ 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: diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 567ecc6..f4676d2 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -115,7 +115,6 @@ 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 = { @@ -174,6 +173,8 @@ 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: