From 2896cf05fb7e3223f19dbdd48dd425d77aadd1eb Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Wed, 25 Jan 2023 22:43:22 +0000
Subject: [PATCH 1/6] v1 working

- from_pretrained support for config
- from_pretrained support for loramodel
- todo: tests
- todo: push_to_hub
---
 src/peft/tuners/lora.py          | 62 +++++++++++++++++++++++++++++---
 src/peft/utils/__init__.py       |  1 +
 src/peft/utils/adapters_utils.py | 17 +++++++++
 src/peft/utils/config.py         | 59 ++++++++++++++++++++++++++----
 4 files changed, 129 insertions(+), 10 deletions(-)
 create mode 100644 src/peft/utils/adapters_utils.py

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index cfd3fd8..5040dc9 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os
 import importlib
 import math
 import warnings
@@ -25,7 +25,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers.pytorch_utils import Conv1D
 
-from ..utils import PeftConfig, PeftType, transpose
+from huggingface_hub import hf_hub_download
+
+from ..utils import PeftConfig, PeftType, transpose, WEIGHTS_NAME, CONFIG_NAME, get_peft_model_state_dict
 
 
 def is_loralib_available():
@@ -85,8 +87,9 @@ class LoraModel(torch.nn.Module):
 
     Example::
 
-        >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>>
-        config = LoraConfig(
+        >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig 
+        >>> from peft import LoraModel, LoraConfig 
+        >>> config = LoraConfig(
             peft_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"],
             lora_dropout=0.01, )
         >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoraModel(config, model)
@@ -167,6 +170,57 @@ class LoraModel(torch.nn.Module):
             config["inference_mode"] = True
         return config
 
+    def save_pretrained(self, save_directory):
+        if os.path.isfile(save_directory):
+            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # save the config
+        self.peft_config.save_pretrained(save_directory)
+
+        for param in self.parameters():
+            param.requires_grad = False  # freeze the model 
+
+        # save only the trainable weights
+        output_state_dict = get_peft_model_state_dict(self)
+        torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
+    
+    @classmethod
+    def from_pretrained(cls, model, lora_id, **kwargs):
+        r"""
+        
+        Args:
+            model (`transformers.PreTrainedModel`): 
+                The model to be adapted.
+            lora_id (`str`):
+                The name of the Lora configuration to use.
+        """
+        # load the config
+        config = LoraConfig.from_pretrained(lora_id)
+        
+        model = cls(config, model)
+
+        # load weights if any
+        if os.path.exists(os.path.join(lora_id, WEIGHTS_NAME)):
+            filename = os.path.join(lora_id, WEIGHTS_NAME)
+        else:
+            try:
+                filename = hf_hub_download(lora_id, WEIGHTS_NAME)
+            except: # noqa
+                raise ValueError(
+                    f"Can't find weights for {lora_id} in {lora_id} or in the Hugging Face Hub. "
+                    f"Please check that the file {WEIGHTS_NAME} is present at {lora_id}."
+                )
+
+
+        adapters_weights = torch.load(filename)
+        # load the weights into the model
+        model.load_state_dict(adapters_weights, strict=False)
+
+        return model
+
+
+
 
 # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
 # and modified to work with PyTorch FSDP
diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py
index 45b0cd5..5a35b1e 100644
--- a/src/peft/utils/__init__.py
+++ b/src/peft/utils/__init__.py
@@ -20,3 +20,4 @@
 from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType
 from .other import _set_trainable, bloom_model_postprocess_past_key_value, shift_tokens_right, transpose
 from .save_and_load import get_peft_model_state_dict, peft_model_load_and_dispatch, set_peft_model_state_dict
+from .adapters_utils import WEIGHTS_NAME, CONFIG_NAME
\ No newline at end of file
diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py
new file mode 100644
index 0000000..9f3761c
--- /dev/null
+++ b/src/peft/utils/adapters_utils.py
@@ -0,0 +1,17 @@
+# coding=utf-8
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+WEIGHTS_NAME = "adapter_model.bin"
+CONFIG_NAME = "adapter_config.json"
\ No newline at end of file
diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
index 17690a0..41bef20 100644
--- a/src/peft/utils/config.py
+++ b/src/peft/utils/config.py
@@ -12,11 +12,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os
+import json
 import enum
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, asdict
 from typing import Optional, Union
 
+from huggingface_hub import hf_hub_download
+
+from .adapters_utils import CONFIG_NAME
+
 
 class PeftType(str, enum.Enum):
     PROMPT_TUNING = "PROMPT_TUNING"
@@ -31,9 +36,53 @@ class TaskType(str, enum.Enum):
     CAUSAL_LM = "CAUSAL_LM"
     TOKEN_CLS = "TOKEN_CLS"
 
-
 @dataclass
-class PeftConfig:
+class PeftConfigMixin(object):
+    @property
+    def __dict__(self):
+        return asdict(self)
+
+    def save_pretrained(self, save_directory):
+        if os.path.isfile(save_directory):
+            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        output_dict = self.__dict__
+        output_path = os.path.join(save_directory, CONFIG_NAME)
+
+        # save it
+        with open(output_path, "w") as writer:
+            writer.write(json.dumps(output_dict))    
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        if os.path.isfile(os.path.join(pretrained_model_name_or_path, CONFIG_NAME)):
+            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+        else:
+            try:
+                config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME)
+            except:
+                raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'")
+        
+        loaded_attributes = cls.from_json_file(config_file)
+        config = cls(**kwargs)
+
+        for key, value in loaded_attributes.items():
+            if hasattr(config, key):
+                setattr(config, key, value)
+
+        return config
+        
+    @classmethod
+    def from_json_file(cls, json_file, **kwargs):
+        with open(json_file, 'r') as file:
+            json_object = json.load(file)
+
+        return json_object
+
+
+class PeftConfig(PeftConfigMixin):
     """
     This is the base configuration class to store the configuration of a :class:`~peft.PeftModel`.
 
@@ -42,13 +91,11 @@ class PeftConfig:
         task_type (Union[[`~peft.utils.config.TaskType`], `str`]): The type of task to perform.
         inference_mode (`bool`, defaults to `False`): Whether to use the Peft model in inference mode.
     """
-
     peft_type: Union[str, PeftType] = field(default=None, metadata={"help": "Peft type"})
     task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"})
     inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"})
 
 
-@dataclass
 class PromptLearningConfig(PeftConfig):
     """
     This is the base configuration class to store the configuration of a Union[[`~peft.PrefixTuning`],

From 2cc7f2cbacf94012f1d287b6a77e0045cb94a667 Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Thu, 26 Jan 2023 10:12:51 +0000
Subject: [PATCH 2/6] add config tests

---
 src/peft/utils/config.py |  5 +++
 tests/test_config.py     | 79 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 tests/test_config.py

diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
index 41bef20..ef985e3 100644
--- a/src/peft/utils/config.py
+++ b/src/peft/utils/config.py
@@ -38,10 +38,15 @@ class TaskType(str, enum.Enum):
 
 @dataclass
 class PeftConfigMixin(object):
+    peft_type: Optional[PeftType] = field(default=None, metadata={"help": "The type of PEFT model."})
+
     @property
     def __dict__(self):
         return asdict(self)
 
+    def to_dict(self):
+        return self.__dict__
+
     def save_pretrained(self, save_directory):
         if os.path.isfile(save_directory):
             raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
diff --git a/tests/test_config.py b/tests/test_config.py
new file mode 100644
index 0000000..7d3df71
--- /dev/null
+++ b/tests/test_config.py
@@ -0,0 +1,79 @@
+import unittest
+import tempfile
+import os
+
+from peft import LoraConfig, PromptEncoderConfig, PrefixTuningConfig, PromptTuningConfig
+
+class PeftConfigMixin:
+    all_config_classes = (
+        LoraConfig,
+        PromptEncoderConfig,
+        PrefixTuningConfig,
+        PromptTuningConfig,
+    )
+
+
+class PeftConfigTester(unittest.TestCase, PeftConfigMixin):
+    def test_methods(self):
+        r"""
+        Test if all configs have the expected methods. Here we test
+        - to_dict
+        - save_pretrained
+        - from_pretrained
+        - from_json_file
+        """
+        # test if all configs have the expected methods
+        for config_class in self.all_config_classes:
+            config = config_class()
+            self.assertTrue(hasattr(config, "to_dict"))
+            self.assertTrue(hasattr(config, "save_pretrained"))
+            self.assertTrue(hasattr(config, "from_pretrained"))
+            self.assertTrue(hasattr(config, "from_json_file"))
+
+
+    def test_save_pretrained(self):
+        r"""
+        Test if the config is correctly saved and loaded using 
+        - save_pretrained
+        """
+        for config_class in self.all_config_classes:
+            config = config_class()
+            with tempfile.TemporaryDirectory() as tmp_dirname:
+                config.save_pretrained(tmp_dirname)
+
+                config_from_pretrained = config_class.from_pretrained(tmp_dirname)
+                self.assertEqual(config.to_dict(), config_from_pretrained.to_dict())
+    
+    def test_from_json_file(self):
+        for config_class in self.all_config_classes:
+            config = config_class()
+            with tempfile.TemporaryDirectory() as tmp_dirname:
+                config.save_pretrained(tmp_dirname)
+
+                config_from_json = config_class.from_json_file(os.path.join(tmp_dirname, "adapter_config.json"))
+                self.assertEqual(config.to_dict(), config_from_json)
+    
+
+    def test_to_dict(self):
+        r"""
+        Test if the config can be correctly converted to a dict using:
+        - to_dict 
+        - __dict__
+        """
+        for config_class in self.all_config_classes:
+            config = config_class()
+            self.assertEqual(config.to_dict(), config.__dict__)
+            self.assertTrue(isinstance(config.to_dict(), dict))
+
+    
+    def test_set_attributes(self):
+        # manually set attributes and check if they are correctly written
+        for config_class in self.all_config_classes:
+            config = config_class(peft_type="test")
+
+            # save pretrained
+            with tempfile.TemporaryDirectory() as tmp_dirname:
+                config.save_pretrained(tmp_dirname)
+
+                config_from_pretrained = config_class.from_pretrained(tmp_dirname)
+                self.assertEqual(config.to_dict(), config_from_pretrained.to_dict())
\ No newline at end of file

From 634f3692d8a186bba5e059c46dbbfe676f31159d Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Thu, 26 Jan 2023 11:17:24 +0000
Subject: [PATCH 3/6] working v1

- push to hub method works
- add tests
- add config super class
- add Lora support for `from_pretrained`
---
 src/peft/tuners/lora.py          |  86 +++++++++++++++++-------
 src/peft/utils/adapters_utils.py |   5 +-
 src/peft/utils/config.py         |  49 ++++++++++++--
 tests/test_config.py             |  18 ++++-
 tests/test_lora.py               | 110 +++++++++++++++++++++++++++++++
 5 files changed, 235 insertions(+), 33 deletions(-)
 create mode 100644 tests/test_lora.py

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 5040dc9..17f75c3 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -27,6 +27,7 @@ from transformers.pytorch_utils import Conv1D
 
 from huggingface_hub import hf_hub_download
 
+from transformers.utils import PushToHubMixin
 from ..utils import PeftConfig, PeftType, transpose, WEIGHTS_NAME, CONFIG_NAME, get_peft_model_state_dict
 
 
@@ -38,6 +39,17 @@ if is_loralib_available():
     import loralib as lora  # noqa: F401
     from loralib import mark_only_lora_as_trainable
 
+MODEL_CARD_TEMPLATE = """---
+license: apache-2.0
+base_model: {base_model}
+tags:
+- peft
+- lora
+---
+# Lora adapters for {model_name}
+
+"""
+
 
 @dataclass
 class LoraConfig(PeftConfig):
@@ -74,7 +86,7 @@ class LoraConfig(PeftConfig):
         self.peft_type = PeftType.LORA
 
 
-class LoraModel(torch.nn.Module):
+class LoraModel(PushToHubMixin, torch.nn.Module):
     """
     Creates Low Rank Adapter (Lora) model from a pretrained transformers model.
 
@@ -150,27 +162,19 @@ class LoraModel(torch.nn.Module):
         if old_module.bias is not None:
             new_module.bias = old_module.bias
 
-    def forward(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+    def save_pretrained(self, save_directory, **kwargs):
+        r"""
+        This function saves the adapter model and the adapter configuration files to a directory, so that it
+        can be re-loaded using the `LoraModel.from_pretrained` class method, and also used by the
+        `LoraModel.push_to_hub` method.
 
-    def __getattr__(self, name: str):
-        """Forward missing attributes to the wrapped module."""
-        try:
-            return super().__getattr__(name)  # defer to nn.Module's logic
-        except AttributeError:
-            return getattr(self.model, name)
-
-    @property
-    def modules_to_save(self):
-        return None
-
-    def get_peft_config_as_dict(self, inference: bool = False):
-        config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(self.peft_config).items()}
-        if inference:
-            config["inference_mode"] = True
-        return config
-
-    def save_pretrained(self, save_directory):
+        Args:   
+            save_directory (`str`):
+                Directory where the adapter model and configuration files will be saved (will be created if it does not
+                exist).
+            **kwargs:
+                Additional keyword arguments passed along to the `push_to_hub` method.
+        """
         if os.path.isfile(save_directory):
             raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
         os.makedirs(save_directory, exist_ok=True)
@@ -184,16 +188,31 @@ class LoraModel(torch.nn.Module):
         # save only the trainable weights
         output_state_dict = get_peft_model_state_dict(self)
         torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
+
+        # save model card
+        if 'name_or_path' in self.model.__dict__:
+            model_name = self.model.__dict__['name_or_path']
+        else:
+            model_name = None
+        model_card_content = MODEL_CARD_TEMPLATE.format(model_name=model_name, base_model=model_name)
+        with open(os.path.join(save_directory, "README.md"), "w", encoding="utf-8") as f:
+            f.write(model_card_content)
     
     @classmethod
     def from_pretrained(cls, model, lora_id, **kwargs):
         r"""
+        Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
         
         Args:
             model (`transformers.PreTrainedModel`): 
-                The model to be adapted.
+                The model to be adapted. The model should be initialized with the `from_pretrained` method.
+                from `transformers` library.
             lora_id (`str`):
-                The name of the Lora configuration to use.
+                The name of the Lora configuration to use. Can be either:
+                    - A string, the `model id` of a Lora configuration hosted inside a model repo on
+                        huggingface Hub
+                    - A path to a directory containing a Lora configuration file saved using the
+                        `save_pretrained` method, e.g., ``./my_lora_config_directory/``.
         """
         # load the config
         config = LoraConfig.from_pretrained(lora_id)
@@ -220,6 +239,27 @@ class LoraModel(torch.nn.Module):
         return model
 
 
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def __getattr__(self, name: str):
+        """Forward missing attributes to the wrapped module."""
+        try:
+            return super().__getattr__(name)  # defer to nn.Module's logic
+        except AttributeError:
+            return getattr(self.model, name)
+
+    @property
+    def modules_to_save(self):
+        return None
+
+    def get_peft_config_as_dict(self, inference: bool = False):
+        config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(self.peft_config).items()}
+        if inference:
+            config["inference_mode"] = True
+        return config
+
+
 
 
 # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py
index 9f3761c..fe906c4 100644
--- a/src/peft/utils/adapters_utils.py
+++ b/src/peft/utils/adapters_utils.py
@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 WEIGHTS_NAME = "adapter_model.bin"
-CONFIG_NAME = "adapter_config.json"
\ No newline at end of file
+CONFIG_NAME = "adapter_config.json"
+
+# TODO: add automapping and superclass here?
\ No newline at end of file
diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
index ef985e3..ad30b32 100644
--- a/src/peft/utils/config.py
+++ b/src/peft/utils/config.py
@@ -18,11 +18,11 @@ import enum
 from dataclasses import dataclass, field, asdict
 from typing import Optional, Union
 
+from transformers.utils import PushToHubMixin
 from huggingface_hub import hf_hub_download
 
 from .adapters_utils import CONFIG_NAME
 
-
 class PeftType(str, enum.Enum):
     PROMPT_TUNING = "PROMPT_TUNING"
     P_TUNING = "P_TUNING"
@@ -36,8 +36,19 @@ class TaskType(str, enum.Enum):
     CAUSAL_LM = "CAUSAL_LM"
     TOKEN_CLS = "TOKEN_CLS"
 
+
 @dataclass
-class PeftConfigMixin(object):
+class PeftConfigMixin(PushToHubMixin):
+    r"""
+    This is the base configuration class for PEFT adapter models. It contains all the methods that 
+    are common to all PEFT adapter models.
+    This class inherits from `transformers.utils.PushToHubMixin` which contains the methods to push your model to the Hub.
+    The method `save_pretrained` will save the configuration of your adapter model in a directory.
+    The method `from_pretrained` will load the configuration of your adapter model from a directory.
+
+    Args:
+        peft_type (Union[[`~peft.utils.config.PeftType`], `str`]): The type of Peft method to use.
+    """
     peft_type: Optional[PeftType] = field(default=None, metadata={"help": "The type of PEFT model."})
 
     @property
@@ -47,7 +58,16 @@ class PeftConfigMixin(object):
     def to_dict(self):
         return self.__dict__
 
-    def save_pretrained(self, save_directory):
+    def save_pretrained(self, save_directory, **kwargs):
+        r"""
+        This method saves the configuration of your adapter model in a directory.
+        
+        Args:
+            save_directory (`str`):
+                The directory where the configuration will be saved.
+            **kwargs:
+                Additional keyword arguments passed along to the `transformers.utils.PushToHubMixin.push_to_hub` method.
+        """
         if os.path.isfile(save_directory):
             raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
 
@@ -58,10 +78,19 @@ class PeftConfigMixin(object):
 
         # save it
         with open(output_path, "w") as writer:
-            writer.write(json.dumps(output_dict))    
+            writer.write(json.dumps(output_dict, indent=2, sort_keys=True))    
     
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        r"""
+        This method loads the configuration of your adapter model from a directory.
+
+        Args:
+            pretrained_model_name_or_path (`str`):
+                The directory or the hub-id where the configuration is saved.
+            **kwargs:
+                Additional keyword arguments passed along to the child class initialization.
+        """
         if os.path.isfile(os.path.join(pretrained_model_name_or_path, CONFIG_NAME)):
             config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
         else:
@@ -71,6 +100,7 @@ class PeftConfigMixin(object):
                 raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'")
         
         loaded_attributes = cls.from_json_file(config_file)
+
         config = cls(**kwargs)
 
         for key, value in loaded_attributes.items():
@@ -80,8 +110,15 @@ class PeftConfigMixin(object):
         return config
         
     @classmethod
-    def from_json_file(cls, json_file, **kwargs):
-        with open(json_file, 'r') as file:
+    def from_json_file(cls, path_json_file, **kwargs):
+        r"""
+        Loads a configuration file from a json file.
+
+        Args:
+            path_json_file (`str`):
+                The path to the json file.
+        """
+        with open(path_json_file, 'r') as file:
             json_object = json.load(file)
 
         return json_object
diff --git a/tests/test_config.py b/tests/test_config.py
index 7d3df71..98881e9 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,10 +1,24 @@
+# coding=utf-8
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import unittest
 import tempfile
 import os
 
 from peft import LoraConfig, PromptEncoderConfig, PrefixTuningConfig, PromptTuningConfig
 
-class PeftConfigMixin:
+class PeftConfigTestMixin:
     all_config_classes = (
         LoraConfig,
         PromptEncoderConfig,
@@ -13,7 +27,7 @@ class PeftConfigMixin:
     )
 
 
-class PeftConfigTester(unittest.TestCase, PeftConfigMixin):
+class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
     def test_methods(self):
         r"""
         Test if all configs have the expected methods. Here we test
diff --git a/tests/test_lora.py b/tests/test_lora.py
new file mode 100644
index 0000000..30a8268
--- /dev/null
+++ b/tests/test_lora.py
@@ -0,0 +1,110 @@
+# coding=utf-8
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import torch
+import tempfile
+import unittest
+
+from peft import LoraModel, LoraConfig, get_peft_model_state_dict
+
+from transformers import AutoModelForCausalLM
+
+class LoraTestMixin:
+    checkpoints_to_test = [
+        "trl-internal-testing/tiny-random-OPTForCausalLM",
+    ]
+
+class LoraTester(unittest.TestCase, LoraTestMixin):
+    r"""
+    Test if the LoraModel behaves as expected. This includes:
+    - test if the model has the expected methods
+    """
+    def test_attributes_lora_model(self):
+        for model_id in self.checkpoints_to_test:
+            model = AutoModelForCausalLM.from_pretrained(model_id)
+            
+            config = LoraConfig(
+                r = 8,
+                lora_alpha=32,
+                target_modules=["q_proj", "v_proj"],
+                lora_dropout=0.05,
+                bias="none",
+            )
+            model = LoraModel(config, model)
+
+            self.assertTrue(hasattr(model, 'save_pretrained'))
+            self.assertTrue(hasattr(model, 'from_pretrained'))
+            self.assertTrue(hasattr(model, 'push_to_hub'))
+    
+    def test_save_pretrained(self):
+        r"""
+        A test to check if `save_pretrained` behaves as expected. This function
+        should only save the state dict of the adapter model and not the state
+        dict of the base model. Hence inside each saved directory you should have:
+
+        - README.md (that contains an entry `base_model`)
+        - adapter_config.json
+        - adapter_model.bin
+        
+        """
+        for model_id in self.checkpoints_to_test:
+            model = AutoModelForCausalLM.from_pretrained(model_id)
+            
+            config = LoraConfig(
+                r = 8,
+                lora_alpha=32,
+                target_modules=["q_proj", "v_proj"],
+                lora_dropout=0.05,
+                bias="none",
+            )
+            model = LoraModel(config, model)
+
+            with tempfile.TemporaryDirectory() as tmp_dirname:
+                model.save_pretrained(tmp_dirname)
+
+                model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id)
+                model_from_pretrained = LoraModel.from_pretrained(model_from_pretrained, tmp_dirname)
+                
+                # check if the state dicts are equal
+                state_dict = get_peft_model_state_dict(model)
+                state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained)
+
+                # check if same keys
+                self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys())
+
+                # check if tensors equal
+                for key in state_dict.keys():
+                    self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key]))
+
+                # check if `README.md` is present
+                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "README.md")))
+                # check if `base_model` attribute is in `README.md`
+                with open(os.path.join(tmp_dirname, "README.md"), "r") as f:
+                    readme = f.read()
+                    self.assertTrue("base_model" in readme)
+
+                # check if `adapter_model.bin` is present
+                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin")))
+
+                # check if `adapter_config.json` is present
+                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json")))
+
+                # check if `pytorch_model.bin` is not present
+                self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin")))
+
+                # check if `config.json` is not present
+                self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))
+
+

From 16182ea972f8097dd429cf014bc8cc6c8bcce5d2 Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Sun, 29 Jan 2023 11:41:38 +0100
Subject: [PATCH 4/6] Apply suggestions from code review

Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
---
 src/peft/tuners/lora.py  | 3 ---
 src/peft/utils/config.py | 1 +
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 7a42572..35d6e4f 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -255,9 +255,6 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
         return model
 
 
-    def forward(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
-
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
         try:
diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
index ad30b32..2a9d264 100644
--- a/src/peft/utils/config.py
+++ b/src/peft/utils/config.py
@@ -124,6 +124,7 @@ class PeftConfigMixin(PushToHubMixin):
         return json_object
 
 
+@dataclass
 class PeftConfig(PeftConfigMixin):
     """
     This is the base configuration class to store the configuration of a :class:`~peft.PeftModel`.

From 22295c427888c18d84758c306f2e7f4d8e3dce7d Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Sun, 29 Jan 2023 10:49:31 +0000
Subject: [PATCH 5/6] adapt from code review

- remove `README`
- inherit from `dataclass`
- add new test
---
 src/peft/tuners/lora.py          | 60 +++++++++-----------------------
 src/peft/utils/__init__.py       |  2 +-
 src/peft/utils/adapters_utils.py |  2 +-
 src/peft/utils/config.py         | 35 +++++++++++--------
 tests/test_config.py             |  5 +++
 tests/test_lora.py               |  7 ----
 6 files changed, 44 insertions(+), 67 deletions(-)

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 35d6e4f..071bc90 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -12,9 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import os
 import importlib
 import math
+import os
 import warnings
 from dataclasses import asdict, dataclass, field
 from enum import Enum
@@ -24,13 +24,12 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers.pytorch_utils import Conv1D
-
+from transformers.utils import PushToHubMixin
 
 import bitsandbytes as bnb
 from huggingface_hub import hf_hub_download
 
-from transformers.utils import PushToHubMixin
-from ..utils import PeftConfig, PeftType, transpose, WEIGHTS_NAME, CONFIG_NAME, get_peft_model_state_dict
+from ..utils import WEIGHTS_NAME, PeftConfig, PeftType, get_peft_model_state_dict, transpose
 
 
 def is_loralib_available():
@@ -41,17 +40,6 @@ if is_loralib_available():
     import loralib as lora  # noqa: F401
     from loralib import mark_only_lora_as_trainable
 
-MODEL_CARD_TEMPLATE = """---
-license: apache-2.0
-base_model: {base_model}
-tags:
-- peft
-- lora
----
-# Lora adapters for {model_name}
-
-"""
-
 
 @dataclass
 class LoraConfig(PeftConfig):
@@ -101,9 +89,8 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
 
     Example::
 
-        >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig 
-        >>> from peft import LoraModel, LoraConfig 
-        >>> config = LoraConfig(
+        >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>>
+        config = LoraConfig(
             peft_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"],
             lora_dropout=0.01, )
         >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoraModel(config, model)
@@ -180,11 +167,11 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
 
     def save_pretrained(self, save_directory, **kwargs):
         r"""
-        This function saves the adapter model and the adapter configuration files to a directory, so that it
-        can be re-loaded using the `LoraModel.from_pretrained` class method, and also used by the
-        `LoraModel.push_to_hub` method.
+        This function saves the adapter model and the adapter configuration files to a directory, so that it can be
+        re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub`
+        method.
 
-        Args:   
+        Args:
             save_directory (`str`):
                 Directory where the adapter model and configuration files will be saved (will be created if it does not
                 exist).
@@ -194,35 +181,26 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
         if os.path.isfile(save_directory):
             raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
         os.makedirs(save_directory, exist_ok=True)
-        
+
         # save the config
         self.peft_config.save_pretrained(save_directory)
 
         for param in self.parameters():
-            param.requires_grad = False  # freeze the model 
+            param.requires_grad = False  # freeze the model
 
         # save only the trainable weights
         output_state_dict = get_peft_model_state_dict(self)
         torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
 
-        # save model card
-        if 'name_or_path' in self.model.__dict__:
-            model_name = self.model.__dict__['name_or_path']
-        else:
-            model_name = None
-        model_card_content = MODEL_CARD_TEMPLATE.format(model_name=model_name, base_model=model_name)
-        with open(os.path.join(save_directory, "README.md"), "w", encoding="utf-8") as f:
-            f.write(model_card_content)
-    
     @classmethod
     def from_pretrained(cls, model, lora_id, **kwargs):
         r"""
         Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
-        
+
         Args:
-            model (`transformers.PreTrainedModel`): 
-                The model to be adapted. The model should be initialized with the `from_pretrained` method.
-                from `transformers` library.
+            model (`transformers.PreTrainedModel`):
+                The model to be adapted. The model should be initialized with the `from_pretrained` method. from
+                `transformers` library.
             lora_id (`str`):
                 The name of the Lora configuration to use. Can be either:
                     - A string, the `model id` of a Lora configuration hosted inside a model repo on
@@ -232,7 +210,7 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
         """
         # load the config
         config = LoraConfig.from_pretrained(lora_id)
-        
+
         model = cls(config, model)
 
         # load weights if any
@@ -241,20 +219,18 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
         else:
             try:
                 filename = hf_hub_download(lora_id, WEIGHTS_NAME)
-            except: # noqa
+            except:  # noqa
                 raise ValueError(
                     f"Can't find weights for {lora_id} in {lora_id} or in the Hugging Face Hub. "
                     f"Please check that the file {WEIGHTS_NAME} is present at {lora_id}."
                 )
 
-
         adapters_weights = torch.load(filename)
         # load the weights into the model
         model.load_state_dict(adapters_weights, strict=False)
 
         return model
 
-
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
         try:
@@ -273,8 +249,6 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
         return config
 
 
-
-
 # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
 # and modified to work with PyTorch FSDP
 
diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py
index 5a35b1e..c418d3d 100644
--- a/src/peft/utils/__init__.py
+++ b/src/peft/utils/__init__.py
@@ -17,7 +17,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .adapters_utils import CONFIG_NAME, WEIGHTS_NAME
 from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType
 from .other import _set_trainable, bloom_model_postprocess_past_key_value, shift_tokens_right, transpose
 from .save_and_load import get_peft_model_state_dict, peft_model_load_and_dispatch, set_peft_model_state_dict
-from .adapters_utils import WEIGHTS_NAME, CONFIG_NAME
\ No newline at end of file
diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py
index fe906c4..f2f8a95 100644
--- a/src/peft/utils/adapters_utils.py
+++ b/src/peft/utils/adapters_utils.py
@@ -15,4 +15,4 @@
 WEIGHTS_NAME = "adapter_model.bin"
 CONFIG_NAME = "adapter_config.json"
 
-# TODO: add automapping and superclass here?
\ No newline at end of file
+# TODO: add automapping and superclass here?
diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
index 2a9d264..f0587fe 100644
--- a/src/peft/utils/config.py
+++ b/src/peft/utils/config.py
@@ -12,17 +12,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import os
-import json
 import enum
-from dataclasses import dataclass, field, asdict
+import json
+import os
+from dataclasses import asdict, dataclass, field
 from typing import Optional, Union
 
 from transformers.utils import PushToHubMixin
+
 from huggingface_hub import hf_hub_download
 
 from .adapters_utils import CONFIG_NAME
 
+
 class PeftType(str, enum.Enum):
     PROMPT_TUNING = "PROMPT_TUNING"
     P_TUNING = "P_TUNING"
@@ -40,11 +42,10 @@ class TaskType(str, enum.Enum):
 @dataclass
 class PeftConfigMixin(PushToHubMixin):
     r"""
-    This is the base configuration class for PEFT adapter models. It contains all the methods that 
-    are common to all PEFT adapter models.
-    This class inherits from `transformers.utils.PushToHubMixin` which contains the methods to push your model to the Hub.
-    The method `save_pretrained` will save the configuration of your adapter model in a directory.
-    The method `from_pretrained` will load the configuration of your adapter model from a directory.
+    This is the base configuration class for PEFT adapter models. It contains all the methods that are common to all
+    PEFT adapter models. This class inherits from `transformers.utils.PushToHubMixin` which contains the methods to
+    push your model to the Hub. The method `save_pretrained` will save the configuration of your adapter model in a
+    directory. The method `from_pretrained` will load the configuration of your adapter model from a directory.
 
     Args:
         peft_type (Union[[`~peft.utils.config.PeftType`], `str`]): The type of Peft method to use.
@@ -61,12 +62,13 @@ class PeftConfigMixin(PushToHubMixin):
     def save_pretrained(self, save_directory, **kwargs):
         r"""
         This method saves the configuration of your adapter model in a directory.
-        
+
         Args:
             save_directory (`str`):
                 The directory where the configuration will be saved.
             **kwargs:
-                Additional keyword arguments passed along to the `transformers.utils.PushToHubMixin.push_to_hub` method.
+                Additional keyword arguments passed along to the `transformers.utils.PushToHubMixin.push_to_hub`
+                method.
         """
         if os.path.isfile(save_directory):
             raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
@@ -78,8 +80,8 @@ class PeftConfigMixin(PushToHubMixin):
 
         # save it
         with open(output_path, "w") as writer:
-            writer.write(json.dumps(output_dict, indent=2, sort_keys=True))    
-    
+            writer.write(json.dumps(output_dict, indent=2, sort_keys=True))
+
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
         r"""
@@ -98,7 +100,7 @@ class PeftConfigMixin(PushToHubMixin):
                 config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME)
             except:
                 raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'")
-        
+
         loaded_attributes = cls.from_json_file(config_file)
 
         config = cls(**kwargs)
@@ -108,7 +110,7 @@ class PeftConfigMixin(PushToHubMixin):
                 setattr(config, key, value)
 
         return config
-        
+
     @classmethod
     def from_json_file(cls, path_json_file, **kwargs):
         r"""
@@ -118,7 +120,7 @@ class PeftConfigMixin(PushToHubMixin):
             path_json_file (`str`):
                 The path to the json file.
         """
-        with open(path_json_file, 'r') as file:
+        with open(path_json_file, "r") as file:
             json_object = json.load(file)
 
         return json_object
@@ -134,11 +136,14 @@ class PeftConfig(PeftConfigMixin):
         task_type (Union[[`~peft.utils.config.TaskType`], `str`]): The type of task to perform.
         inference_mode (`bool`, defaults to `False`): Whether to use the Peft model in inference mode.
     """
+
+    base_model_name_or_path: str = field(default=None, metadata={"help": "The name of the base model to use."})
     peft_type: Union[str, PeftType] = field(default=None, metadata={"help": "Peft type"})
     task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"})
     inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"})
 
 
+@dataclass
 class PromptLearningConfig(PeftConfig):
     """
     This is the base configuration class to store the configuration of a Union[[`~peft.PrefixTuning`],
diff --git a/tests/test_config.py b/tests/test_config.py
index 98881e9..5a87e44 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -43,6 +43,11 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
             self.assertTrue(hasattr(config, "save_pretrained"))
             self.assertTrue(hasattr(config, "from_pretrained"))
             self.assertTrue(hasattr(config, "from_json_file"))
+    
+    def test_task_type(self):
+        for config_class in self.all_config_classes:
+            # assert this will not fail
+            _ = config_class(task_type="test")
 
 
     def test_save_pretrained(self):
diff --git a/tests/test_lora.py b/tests/test_lora.py
index 30a8268..c69a32a 100644
--- a/tests/test_lora.py
+++ b/tests/test_lora.py
@@ -88,13 +88,6 @@ class LoraTester(unittest.TestCase, LoraTestMixin):
                 for key in state_dict.keys():
                     self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key]))
 
-                # check if `README.md` is present
-                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "README.md")))
-                # check if `base_model` attribute is in `README.md`
-                with open(os.path.join(tmp_dirname, "README.md"), "r") as f:
-                    readme = f.read()
-                    self.assertTrue("base_model" in readme)
-
                 # check if `adapter_model.bin` is present
                 self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin")))
 

From 6c9534e660de73fa44abaefbf1f4f5b0df818fbb Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Sun, 29 Jan 2023 11:18:31 +0000
Subject: [PATCH 6/6] adapt for other models

---
 src/peft/mapping.py     |   2 +
 src/peft/peft_model.py  |  84 ++++++++++++++++++++++++++++-
 src/peft/tuners/lora.py |  73 +------------------------
 tests/test_lora.py      | 114 ++++++++++++++++++++++++----------------
 4 files changed, 154 insertions(+), 119 deletions(-)

diff --git a/src/peft/mapping.py b/src/peft/mapping.py
index 66c07ff..18e292b 100644
--- a/src/peft/mapping.py
+++ b/src/peft/mapping.py
@@ -138,4 +138,6 @@ def get_peft_model(model, peft_config):
     else:
         peft_config = _prepare_lora_config(peft_config, model_config)
 
+    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
+
     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 d43baf8..f3e01c1 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -14,18 +14,31 @@
 # limitations under the License.
 
 import inspect
+import os
 import warnings
 
 import torch
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput
+from transformers.utils import PushToHubMixin
+
+from huggingface_hub import hf_hub_download
 
 from .tuners import LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder
-from .utils import PeftConfig, PeftType, TaskType, _set_trainable, shift_tokens_right
+from .utils import (
+    WEIGHTS_NAME,
+    PeftConfig,
+    PeftType,
+    TaskType,
+    _set_trainable,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+    shift_tokens_right,
+)
 
 
-class PeftModel(torch.nn.Module):
+class PeftModel(PushToHubMixin, torch.nn.Module):
     """
     Parameter-Efficient Fine-Tuning Model. Base model encompassing various Peft methods.
 
@@ -61,6 +74,73 @@ class PeftModel(torch.nn.Module):
             self.base_model = LoraModel(peft_config, model)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+    def save_pretrained(self, save_directory, **kwargs):
+        r"""
+        Args:
+        This function saves the adapter model and the adapter configuration files to a directory, so that it can be
+        re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub`
+        method.
+            save_directory (`str`):
+                Directory where the adapter model and configuration files will be saved (will be created if it does not
+                exist).
+            **kwargs:
+                Additional keyword arguments passed along to the `push_to_hub` method.
+        """
+        if os.path.isfile(save_directory):
+            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save the config
+        if self.peft_config.base_model_name_or_path is None:
+            self.peft_config.base_model_name_or_path = self.base_model.__dict__.get("name_or_path", None)
+        self.peft_config.inference_mode = True
+        self.peft_config.save_pretrained(save_directory)
+
+        for param in self.parameters():
+            param.requires_grad = False  # freeze the model
+
+        # save only the trainable weights
+        output_state_dict = get_peft_model_state_dict(self, kwargs.get("state_dict", None))
+        torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
+
+    @classmethod
+    def from_pretrained(cls, model, model_id, **kwargs):
+        r"""
+        Args:
+        Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
+            model (`transformers.PreTrainedModel`):
+                The model to be adapted. The model should be initialized with the `from_pretrained` method. from
+                `transformers` library.
+            model_id (`str`):
+                The name of the Lora configuration to use. Can be either:
+                    - A string, the `model id` of a Lora configuration hosted inside a model repo on
+                        huggingface Hub
+                    - A path to a directory containing a Lora configuration file saved using the
+                        `save_pretrained` method, e.g., ``./my_lora_config_directory/``.
+        """
+        from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING
+
+        # load the config
+        config = PEFT_TYPE_TO_CONFIG_MAPPING[PeftConfig.from_pretrained(model_id).peft_type].from_pretrained(model_id)
+
+        model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config)
+
+        # load weights if any
+        if os.path.exists(os.path.join(model_id, WEIGHTS_NAME)):
+            filename = os.path.join(model_id, WEIGHTS_NAME)
+        else:
+            try:
+                filename = hf_hub_download(model_id, WEIGHTS_NAME)
+            except:  # noqa
+                raise ValueError(
+                    f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. "
+                    f"Please check that the file {WEIGHTS_NAME} is present at {model_id}."
+                )
+
+        adapters_weights = torch.load(filename)
+        # load the weights into the model
+        return set_peft_model_state_dict(model, adapters_weights)
+
     def _setup_prompt_encoder(self):
         num_transformer_submodules = 0
         transformer_backbone = None
diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 071bc90..a1eb3b1 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -14,7 +14,6 @@
 # limitations under the License.
 import importlib
 import math
-import os
 import warnings
 from dataclasses import asdict, dataclass, field
 from enum import Enum
@@ -24,12 +23,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers.pytorch_utils import Conv1D
-from transformers.utils import PushToHubMixin
 
 import bitsandbytes as bnb
-from huggingface_hub import hf_hub_download
 
-from ..utils import WEIGHTS_NAME, PeftConfig, PeftType, get_peft_model_state_dict, transpose
+from ..utils import PeftConfig, PeftType, transpose
 
 
 def is_loralib_available():
@@ -76,7 +73,7 @@ class LoraConfig(PeftConfig):
         self.peft_type = PeftType.LORA
 
 
-class LoraModel(PushToHubMixin, torch.nn.Module):
+class LoraModel(torch.nn.Module):
     """
     Creates Low Rank Adapter (Lora) model from a pretrained transformers model.
 
@@ -165,72 +162,6 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
             new_module.state = old_module.state
             new_module.to(old_module.weight.device)
 
-    def save_pretrained(self, save_directory, **kwargs):
-        r"""
-        This function saves the adapter model and the adapter configuration files to a directory, so that it can be
-        re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub`
-        method.
-
-        Args:
-            save_directory (`str`):
-                Directory where the adapter model and configuration files will be saved (will be created if it does not
-                exist).
-            **kwargs:
-                Additional keyword arguments passed along to the `push_to_hub` method.
-        """
-        if os.path.isfile(save_directory):
-            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
-        os.makedirs(save_directory, exist_ok=True)
-
-        # save the config
-        self.peft_config.save_pretrained(save_directory)
-
-        for param in self.parameters():
-            param.requires_grad = False  # freeze the model
-
-        # save only the trainable weights
-        output_state_dict = get_peft_model_state_dict(self)
-        torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
-
-    @classmethod
-    def from_pretrained(cls, model, lora_id, **kwargs):
-        r"""
-        Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
-
-        Args:
-            model (`transformers.PreTrainedModel`):
-                The model to be adapted. The model should be initialized with the `from_pretrained` method. from
-                `transformers` library.
-            lora_id (`str`):
-                The name of the Lora configuration to use. Can be either:
-                    - A string, the `model id` of a Lora configuration hosted inside a model repo on
-                        huggingface Hub
-                    - A path to a directory containing a Lora configuration file saved using the
-                        `save_pretrained` method, e.g., ``./my_lora_config_directory/``.
-        """
-        # load the config
-        config = LoraConfig.from_pretrained(lora_id)
-
-        model = cls(config, model)
-
-        # load weights if any
-        if os.path.exists(os.path.join(lora_id, WEIGHTS_NAME)):
-            filename = os.path.join(lora_id, WEIGHTS_NAME)
-        else:
-            try:
-                filename = hf_hub_download(lora_id, WEIGHTS_NAME)
-            except:  # noqa
-                raise ValueError(
-                    f"Can't find weights for {lora_id} in {lora_id} or in the Hugging Face Hub. "
-                    f"Please check that the file {WEIGHTS_NAME} is present at {lora_id}."
-                )
-
-        adapters_weights = torch.load(filename)
-        # load the weights into the model
-        model.load_state_dict(adapters_weights, strict=False)
-
-        return model
-
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
         try:
diff --git a/tests/test_lora.py b/tests/test_lora.py
index c69a32a..17b710a 100644
--- a/tests/test_lora.py
+++ b/tests/test_lora.py
@@ -17,36 +17,62 @@ import torch
 import tempfile
 import unittest
 
-from peft import LoraModel, LoraConfig, get_peft_model_state_dict
+from peft import PeftConfig, PeftModel, LoraConfig, get_peft_model_state_dict, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
 
 from transformers import AutoModelForCausalLM
 
 class LoraTestMixin:
     checkpoints_to_test = [
-        "trl-internal-testing/tiny-random-OPTForCausalLM",
+        "hf-internal-testing/tiny-random-OPTForCausalLM",
     ]
+    config_classes = (
+        LoraConfig,
+        # PrefixTuningConfig,
+        # PromptEncoderConfig,
+        # PromptTuningConfig,
+    )
+    config_kwargs = (
+        dict(
+            r = 8,
+            lora_alpha=32,
+            target_modules=["q_proj", "v_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+        ),
+        # dict(
+        #     encoder_hidden_size=32,
+        #     task_type="CAUSAL_LM",
+        # ),
+        # dict(
+        #     encoder_hidden_size=32,
+        #     task_type="CAUSAL_LM",
+        # ),
+        # dict(
+        #     task_type="CAUSAL_LM",
+        # )
 
-class LoraTester(unittest.TestCase, LoraTestMixin):
+    )
+
+class PeftModelTester(unittest.TestCase, LoraTestMixin):
     r"""
-    Test if the LoraModel behaves as expected. This includes:
+    Test if the PeftModel behaves as expected. This includes:
     - test if the model has the expected methods
     """
     def test_attributes_lora_model(self):
         for model_id in self.checkpoints_to_test:
             model = AutoModelForCausalLM.from_pretrained(model_id)
             
-            config = LoraConfig(
-                r = 8,
-                lora_alpha=32,
-                target_modules=["q_proj", "v_proj"],
-                lora_dropout=0.05,
-                bias="none",
-            )
-            model = LoraModel(config, model)
+            for i, config_cls in enumerate(self.config_classes):
+                config = config_cls(
+                    base_model_name_or_path=model_id,
+                    **self.config_kwargs[i],
+                )
+                model = PeftModel(model, config)
 
-            self.assertTrue(hasattr(model, 'save_pretrained'))
-            self.assertTrue(hasattr(model, 'from_pretrained'))
-            self.assertTrue(hasattr(model, 'push_to_hub'))
+                self.assertTrue(hasattr(model, 'save_pretrained'))
+                self.assertTrue(hasattr(model, 'from_pretrained'))
+                self.assertTrue(hasattr(model, 'push_to_hub'))
     
     def test_save_pretrained(self):
         r"""
@@ -62,42 +88,38 @@ class LoraTester(unittest.TestCase, LoraTestMixin):
         for model_id in self.checkpoints_to_test:
             model = AutoModelForCausalLM.from_pretrained(model_id)
             
-            config = LoraConfig(
-                r = 8,
-                lora_alpha=32,
-                target_modules=["q_proj", "v_proj"],
-                lora_dropout=0.05,
-                bias="none",
-            )
-            model = LoraModel(config, model)
+            for i, config_cls in enumerate(self.config_classes):
+                config = config_cls(
+                    base_model_name_or_path=model_id,
+                    **self.config_kwargs[i],
+                )
+                model = PeftModel(model, config)
 
-            with tempfile.TemporaryDirectory() as tmp_dirname:
-                model.save_pretrained(tmp_dirname)
+                with tempfile.TemporaryDirectory() as tmp_dirname:
+                    model.save_pretrained(tmp_dirname)
 
-                model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id)
-                model_from_pretrained = LoraModel.from_pretrained(model_from_pretrained, tmp_dirname)
-                
-                # check if the state dicts are equal
-                state_dict = get_peft_model_state_dict(model)
-                state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained)
+                    model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id)
+                    model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname)
+                    
+                    # check if the state dicts are equal
+                    state_dict = get_peft_model_state_dict(model)
+                    state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained)
 
-                # check if same keys
-                self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys())
+                    # check if same keys
+                    self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys())
 
-                # check if tensors equal
-                for key in state_dict.keys():
-                    self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key]))
+                    # check if tensors equal
+                    for key in state_dict.keys():
+                        self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key]))
 
-                # check if `adapter_model.bin` is present
-                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin")))
+                    # check if `adapter_model.bin` is present
+                    self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin")))
 
-                # check if `adapter_config.json` is present
-                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json")))
-
-                # check if `pytorch_model.bin` is not present
-                self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin")))
-
-                # check if `config.json` is not present
-                self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))
+                    # check if `adapter_config.json` is present
+                    self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json")))
 
+                    # check if `pytorch_model.bin` is not present
+                    self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin")))
 
+                    # check if `config.json` is not present
+                    self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))
\ No newline at end of file