File size: 4,898 Bytes
96fe658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import os
import shutil
import tempfile
import unittest
import torch
from modelscope import AutoModel, Preprocessor
from peft.utils import WEIGHTS_NAME
from transformers import PreTrainedModel
from swift import LoRAConfig, Swift
from swift.tuners import NEFTuneConfig
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, 'pytorch_model.bin')
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, 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', WEIGHTS_NAME)
self.assertTrue(os.path.isfile(bin_file))
bin_file = os.path.join(self.tmp_dir, 'c1', 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()))
|