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

[Variant] Make sure variant files are not incorrectly deleted (#21562)

Browse Source 
* [Variant] Make sure variant files are not incorrectly deleted

* Apply suggestions from code review

* fix
This commit is contained in:
Patrick von Platen
2023-02-10 16:44:51 +02:00
committed by GitHub
parent 51c3f42d8e
commit b20147a3c8
2 changed files with 27 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
tests/test_modeling_common.py
View File

@@ -3119,6 +3119,27 @@ class ModelUtilsTest(TestCasePlus):
)
self.assertIsNotNone(model)
def test_checkpoint_variant_save_load(self):
with tempfile.TemporaryDirectory() as tmp_dir:
model = BertModel.from_pretrained(
"hf-internal-testing/tiny-random-bert-variant", cache_dir=tmp_dir, variant="v2"
)
weights_name = ".".join(WEIGHTS_NAME.split(".")[:-1] + ["v2"] + ["bin"])
model.save_pretrained(tmp_dir, variant="v2")
# saving will create a variant checkpoint
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, weights_name)))
model.save_pretrained(tmp_dir)
# saving shouldn't delete variant checkpoints
weights_name = ".".join(WEIGHTS_NAME.split(".")[:-1] + ["v2"] + ["bin"])
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, weights_name)))
# there should be a normal checkpoint
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, WEIGHTS_NAME)))
self.assertIsNotNone(model)
@require_accelerate
def test_from_pretrained_low_cpu_mem_usage_functional(self):
# test that we can use `from_pretrained(..., low_cpu_mem_usage=True)` with normal and