wehub-resource-sync
108 lines
4.8 KiB
Python
108 lines
4.8 KiB
Python
import os
|
|
import shutil
|
|
import tempfile
|
|
import torch
|
|
import unittest
|
|
from modelscope import AutoModel, Preprocessor
|
|
from peft.utils import SAFETENSORS_WEIGHTS_NAME
|
|
from transformers import PreTrainedModel
|
|
|
|
from swift.tuners import LoRAConfig, NEFTuneConfig, Swift
|
|
|
|
|
|
class TestNEFT(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
|
|
self.tmp_dir = tempfile.TemporaryDirectory().name
|
|
if not os.path.exists(self.tmp_dir):
|
|
os.makedirs(self.tmp_dir)
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.tmp_dir)
|
|
super().tearDown()
|
|
|
|
def test_neft(self):
|
|
model = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
|
|
preprocessor = Preprocessor.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
|
|
inputs = preprocessor('how are you')
|
|
config = NEFTuneConfig()
|
|
|
|
t1 = model.embeddings.word_embeddings(inputs['input_ids'])
|
|
model = Swift.prepare_model(model, config)
|
|
model.train()
|
|
t2 = model.embeddings.word_embeddings(inputs['input_ids'])
|
|
model.deactivate_adapter('default')
|
|
t3 = model.embeddings.word_embeddings(inputs['input_ids'])
|
|
self.assertTrue(torch.allclose(t1, t3))
|
|
self.assertFalse(torch.allclose(t1, t2))
|
|
model.save_pretrained(self.tmp_dir)
|
|
bin_file = os.path.join(self.tmp_dir, 'model.safetensors')
|
|
self.assertTrue(os.path.isfile(bin_file))
|
|
model2 = AutoModel.from_pretrained(self.tmp_dir)
|
|
|
|
state_dict = model.state_dict()
|
|
state_dict2 = model2.state_dict()
|
|
self.assertTrue(len(state_dict) > 0)
|
|
for key in state_dict:
|
|
self.assertTrue(key in state_dict2)
|
|
self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu()))
|
|
|
|
shutil.rmtree(self.tmp_dir)
|
|
PreTrainedModel.origin_save_pretrained = PreTrainedModel.save_pretrained
|
|
delattr(PreTrainedModel, 'save_pretrained')
|
|
model.save_pretrained(self.tmp_dir)
|
|
bin_file = os.path.join(self.tmp_dir, SAFETENSORS_WEIGHTS_NAME)
|
|
self.assertTrue(os.path.isfile(bin_file))
|
|
model_new = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
|
|
model_new_2 = Swift.from_pretrained(model_new, self.tmp_dir)
|
|
|
|
state_dict = model.state_dict()
|
|
state_dict2 = model_new_2.state_dict()
|
|
self.assertTrue(len(state_dict) > 0)
|
|
for key in state_dict:
|
|
self.assertTrue(key in state_dict2)
|
|
self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu()))
|
|
PreTrainedModel.save_pretrained = PreTrainedModel.origin_save_pretrained
|
|
|
|
def test_neft_lora(self):
|
|
model = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
|
|
preprocessor = Preprocessor.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
|
|
inputs = preprocessor('how are you')
|
|
config = NEFTuneConfig()
|
|
config2 = LoRAConfig(target_modules=['query', 'key', 'value'])
|
|
|
|
t1 = model.embeddings.word_embeddings(inputs['input_ids'])
|
|
model = Swift.prepare_model(model, {'c1': config, 'c2': config2})
|
|
model.train()
|
|
t2 = model.embeddings.word_embeddings(inputs['input_ids'])
|
|
model.deactivate_adapter('c1')
|
|
t3 = model.embeddings.word_embeddings(inputs['input_ids'])
|
|
self.assertTrue(torch.allclose(t1, t3))
|
|
self.assertFalse(torch.allclose(t1, t2))
|
|
model.save_pretrained(self.tmp_dir)
|
|
bin_file = os.path.join(self.tmp_dir, 'c2', SAFETENSORS_WEIGHTS_NAME)
|
|
self.assertTrue(os.path.isfile(bin_file))
|
|
bin_file = os.path.join(self.tmp_dir, 'c1', SAFETENSORS_WEIGHTS_NAME)
|
|
self.assertTrue(not os.path.isfile(bin_file))
|
|
model_new = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
|
|
t1 = model_new.embeddings.word_embeddings(inputs['input_ids'])
|
|
model_new = Swift.from_pretrained(model_new, self.tmp_dir)
|
|
model_new.train()
|
|
t2 = model_new.embeddings.word_embeddings(inputs['input_ids'])
|
|
model_new.eval()
|
|
t4 = model_new.embeddings.word_embeddings(inputs['input_ids'])
|
|
model_new.train()
|
|
model_new.deactivate_adapter('c1')
|
|
t3 = model_new.embeddings.word_embeddings(inputs['input_ids'])
|
|
self.assertTrue(torch.allclose(t1, t3))
|
|
self.assertTrue(torch.allclose(t1, t4))
|
|
self.assertFalse(torch.allclose(t1, t2))
|
|
|
|
state_dict = model.state_dict()
|
|
state_dict2 = model_new.state_dict()
|
|
self.assertTrue(len(state_dict) > 0 and all(['lora' in key for key in state_dict.keys()]))
|
|
for key in state_dict:
|
|
self.assertTrue(key in state_dict2)
|
|
self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu()))
|