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:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user