mirror of
https://github.com/wassname/peft.git
synced 2026-09-25 13:50:20 +08:00
younesbelkada
634f3692d8
- push to hub method works - add tests - add config super class - add Lora support for `from_pretrained`
111 lines
4.3 KiB
Python
111 lines
4.3 KiB
Python
# 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")))
|
|
|
|
|