adapt from code review

- remove `README`
- inherit from `dataclass`
- add new test
This commit is contained in:
younesbelkada
2023-01-29 10:49:31 +00:00
parent 16182ea972
commit 22295c4278
6 changed files with 44 additions and 67 deletions
+17 -43
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -15,4 +15,4 @@
WEIGHTS_NAME = "adapter_model.bin"
CONFIG_NAME = "adapter_config.json"
# TODO: add automapping and superclass here?
# TODO: add automapping and superclass here?
+20 -15
View File
@@ -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`],
+5
View File
@@ -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):
-7
View File
@@ -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")))