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

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files

8
src/transformers/modeling_utils.py
View File

@@ -3270,8 +3270,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
) )
else: else:
raise EnvironmentError( raise EnvironmentError(
f"Error no file named {_add_variant(WEIGHTS_NAME, variant)}, {TF2_WEIGHTS_NAME}," f"Error no file named {_add_variant(WEIGHTS_NAME, variant)}, {_add_variant(SAFE_WEIGHTS_NAME, variant)},"
f" {TF_WEIGHTS_NAME + '.index'} or {FLAX_WEIGHTS_NAME} found in directory" f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME + '.index'} or {FLAX_WEIGHTS_NAME} found in directory"
f" {pretrained_model_name_or_path}." f" {pretrained_model_name_or_path}."
) )
elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)): elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
@@ -3417,8 +3417,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
else: else:
raise EnvironmentError( raise EnvironmentError(
f"{pretrained_model_name_or_path} does not appear to have a file named" f"{pretrained_model_name_or_path} does not appear to have a file named"
f" {_add_variant(WEIGHTS_NAME, variant)}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or" f" {_add_variant(WEIGHTS_NAME, variant)}, {_add_variant(SAFE_WEIGHTS_NAME, variant)},"
f" {FLAX_WEIGHTS_NAME}." f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}."
) )
except EnvironmentError: except EnvironmentError:
# Raise any environment error raise by `cached_file`. It will have a helpful error message adapted # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted

20
tests/test_modeling_utils.py
View File

@@ -1001,6 +1001,26 @@ class ModelUtilsTest(TestCasePlus):
self.assertTrue(any(f.endswith("safetensors") for f in all_downloaded_files)) self.assertTrue(any(f.endswith("safetensors") for f in all_downloaded_files))
self.assertFalse(any(f.endswith("bin") 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 @require_safetensors
def test_safetensors_save_and_load(self): def test_safetensors_save_and_load(self):
model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert") model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")