diff --git a/Makefile b/Makefile index 532a640..8d8148e 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ .PHONY: quality style test docs -check_dirs := src examples +check_dirs := src tests examples # Check that source code meets quality standards @@ -9,11 +9,11 @@ quality: black --check $(check_dirs) isort --check-only $(check_dirs) flake8 $(check_dirs) - doc-builder style src --max_len 119 --check_only + doc-builder style src tests --max_len 119 --check_only # Format source code automatically and check is there are any problems left that need manual fixing style: black $(check_dirs) isort $(check_dirs) - doc-builder style src --max_len 119 + doc-builder style src tests --max_len 119 \ No newline at end of file diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index d2c1cd2..9b6962c 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -50,7 +50,10 @@ def get_peft_model_state_dict(model, state_dict=None): raise NotImplementedError else: to_return = {} - prompt_embeddings = model.get_prompt_embedding_to_save() + if model.peft_config.inference_mode: + prompt_embeddings = model.prompt_encoder.embedding.weight + else: + prompt_embeddings = model.get_prompt_embedding_to_save() to_return["prompt_embeddings"] = prompt_embeddings if model.modules_to_save is not None: for key, value in state_dict.items(): diff --git a/tests/test_config.py b/tests/test_config.py index 5a87e44..34fdc44 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -12,11 +12,12 @@ # 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 +import tempfile +import unittest + +from peft import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig -from peft import LoraConfig, PromptEncoderConfig, PrefixTuningConfig, PromptTuningConfig class PeftConfigTestMixin: all_config_classes = ( @@ -43,16 +44,15 @@ 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): r""" - Test if the config is correctly saved and loaded using + Test if the config is correctly saved and loaded using - save_pretrained """ for config_class in self.all_config_classes: @@ -62,7 +62,7 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin): 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() @@ -71,12 +71,11 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin): 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 + - to_dict - __dict__ """ for config_class in self.all_config_classes: @@ -84,7 +83,6 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin): 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: @@ -95,4 +93,4 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin): 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 + self.assertEqual(config.to_dict(), config_from_pretrained.to_dict()) diff --git a/tests/test_lora.py b/tests/test_save_and_load.py similarity index 70% rename from tests/test_lora.py rename to tests/test_save_and_load.py index 17b710a..e17f558 100644 --- a/tests/test_lora.py +++ b/tests/test_save_and_load.py @@ -13,94 +13,105 @@ # 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 PeftConfig, PeftModel, LoraConfig, get_peft_model_state_dict, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig - +import torch from transformers import AutoModelForCausalLM -class LoraTestMixin: +from peft import ( + LoraConfig, + PeftModel, + PrefixTuningConfig, + PromptEncoderConfig, + PromptTuningConfig, + get_peft_model, + get_peft_model_state_dict, +) + + +class PeftTestMixin: checkpoints_to_test = [ "hf-internal-testing/tiny-random-OPTForCausalLM", ] config_classes = ( LoraConfig, - # PrefixTuningConfig, - # PromptEncoderConfig, - # PromptTuningConfig, + PrefixTuningConfig, + PromptEncoderConfig, + PromptTuningConfig, ) config_kwargs = ( dict( - r = 8, + 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", - # ) - + dict( + num_virtual_tokens=10, + task_type="CAUSAL_LM", + ), + dict( + num_virtual_tokens=10, + encoder_hidden_size=32, + task_type="CAUSAL_LM", + ), + dict( + num_virtual_tokens=10, + task_type="CAUSAL_LM", + ), ) -class PeftModelTester(unittest.TestCase, LoraTestMixin): + +class PeftModelTester(unittest.TestCase, PeftTestMixin): r""" Test if the PeftModel behaves as expected. This includes: - test if the model has the expected methods """ - def test_attributes_lora_model(self): + + def test_attributes_model(self): for model_id in self.checkpoints_to_test: - model = AutoModelForCausalLM.from_pretrained(model_id) - for i, config_cls in enumerate(self.config_classes): + model = AutoModelForCausalLM.from_pretrained(model_id) config = config_cls( base_model_name_or_path=model_id, **self.config_kwargs[i], ) - model = PeftModel(model, config) + model = get_peft_model(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""" - 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: + 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) - for i, config_cls in enumerate(self.config_classes): + model = AutoModelForCausalLM.from_pretrained(model_id) config = config_cls( base_model_name_or_path=model_id, **self.config_kwargs[i], ) - model = PeftModel(model, config) + model = get_peft_model(model, config) + model.to(model.device) with tempfile.TemporaryDirectory() as tmp_dirname: model.save_pretrained(tmp_dirname) model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id) model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname) - + model_from_pretrained.to(model.device) + # 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) @@ -122,4 +133,4 @@ class PeftModelTester(unittest.TestCase, LoraTestMixin): 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 + self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))