From 22295c427888c18d84758c306f2e7f4d8e3dce7d Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Sun, 29 Jan 2023 10:49:31 +0000 Subject: [PATCH] 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")))