SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Make sure dynamic objects can be saved and reloaded (#21008)

Browse Source 
* Make sure dynamic objects can be saved and reloaded

* Remove processor test
This commit is contained in:
Sylvain Gugger
2023-01-05 07:30:25 -05:00
committed by GitHub
parent bf82c9b74f
commit 12313838d3
12 changed files with 64 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tests/models/auto/test_modeling_auto.py
View File

@@ -276,10 +276,28 @@ class AutoModelTest(unittest.TestCase):
model = AutoModel.from_pretrained("hf-internal-testing/test_dynamic_model", trust_remote_code=True)
self.assertEqual(model.__class__.__name__, "NewModel")
# Test model can be reloaded.
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir)
reloaded_model = AutoModel.from_pretrained(tmp_dir, trust_remote_code=True)
self.assertEqual(reloaded_model.__class__.__name__, "NewModel")
for p1, p2 in zip(model.parameters(), reloaded_model.parameters()):
self.assertTrue(torch.equal(p1, p2))
# This one uses a relative import to a util file, this checks it is downloaded and used properly.
model = AutoModel.from_pretrained("hf-internal-testing/test_dynamic_model_with_util", trust_remote_code=True)
self.assertEqual(model.__class__.__name__, "NewModel")
# Test model can be reloaded.
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir)
reloaded_model = AutoModel.from_pretrained(tmp_dir, trust_remote_code=True)
self.assertEqual(reloaded_model.__class__.__name__, "NewModel")
for p1, p2 in zip(model.parameters(), reloaded_model.parameters()):
self.assertTrue(torch.equal(p1, p2))
def test_new_model_registration(self):
AutoConfig.register("custom", CustomConfig)