Add safetensors to model not found error msg for default use_safetensors value (#30602)

* add safetensors to model not found error for default use_safetensors=None case

* format code w/ ruff

* fix assert true typo
This commit is contained in:
David Xue
2024-05-07 12:55:27 -04:00
committed by GitHub
parent 884e3b1c53
commit cf7bed9832
2 changed files with 24 additions and 4 deletions

View File

@@ -1001,6 +1001,26 @@ class ModelUtilsTest(TestCasePlus):
self.assertTrue(any(f.endswith("safetensors") for f in all_downloaded_files))
self.assertFalse(any(f.endswith("bin") for f in all_downloaded_files))
# test no model file found when use_safetensors=None (default when safetensors package available)
with self.assertRaises(OSError) as missing_model_file_error:
BertModel.from_pretrained("hf-internal-testing/config-no-model")
self.assertTrue(
"does not appear to have a file named pytorch_model.bin, model.safetensors,"
in str(missing_model_file_error.exception)
)
with self.assertRaises(OSError) as missing_model_file_error:
with tempfile.TemporaryDirectory() as tmp_dir:
with open(os.path.join(tmp_dir, "config.json"), "w") as f:
f.write("{}")
f.close()
BertModel.from_pretrained(tmp_dir)
self.assertTrue(
"Error no file named pytorch_model.bin, model.safetensors" in str(missing_model_file_error.exception)
)
@require_safetensors
def test_safetensors_save_and_load(self):
model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")