Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix ModelHubMixin: save config only if doesn't exist #2105

Merged
merged 2 commits into from
Mar 19, 2024
Merged

Fix ModelHubMixin: save config only if doesn't exist #2105

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8c40a02
Fix ModelHubMixin: save config only if doesn't exist
Wauplin Mar 11, 2024
70bbf93
sutff
Wauplin Mar 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions src/huggingface_hub/hub_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -248,13 +248,16 @@ def save_pretrained(
# save model weights/files (framework-specific)
self._save_pretrained(save_directory)

# save config (if provided)
# save config (if provided and if not serialized yet in `_save_pretrained`)
if config is None:
config = self.config
if config is not None:
if is_dataclass(config):
config = asdict(config) # type: ignore[arg-type]
(save_directory / CONFIG_NAME).write_text(json.dumps(config, sort_keys=True, indent=2))
config_path = save_directory / CONFIG_NAME
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This might not work if the user wants to save a file with a name other than the value of the constant CONFIG_NAME, e.g. tokenizer.json. A library could potentially use a different filename by default, and implement a filename kwarg passed in push_to_hub_kwargs.
In such case, this can be bypassed by reassigning CONFIG_NAME in save_pretrained, but there might be a better way to handle this to make sure it doesn't interfere with anything else.

Maybe ModelHubMixin could implement a protected _config_filename attribute whose value would default to CONFIG_NAME, that libraries could freely reassign?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the concern and suggestion @Natooz!

Actually I would prefer not to allow different config files, at least until we witness a real case in which it would be beneficial. config.json is a special filename on the Hub as it is automatically used to count model downloads. If the library author wants 2 configs files -1 handled by the library, 1 handled by huggingface_hub-, then it'd be best if the library-specific config file is named differently (keras_config.json, tokenizer_config.json, etc.).

if not config_path.exists():
config_str = json.dumps(config, sort_keys=True, indent=2)
config_path.write_text(config_str)

# save model card
model_card_path = save_directory / "README.md"
Expand Down
17 changes: 17 additions & 0 deletions tests/test_hub_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,15 @@ def __init__(self, **kwargs):
pass


class DummyModelSavingConfig(ModelHubMixin):
def _save_pretrained(self, save_directory: Path) -> None:
"""Implementation that uses `config.json` to serialize the config.

This file must not be overwritten by the default config saved by `ModelHubMixin`.
"""
(save_directory / "config.json").write_text(json.dumps({"custom_config": "custom_config"}))


@pytest.mark.usefixtures("fx_cache_dir")
class HubMixinTest(unittest.TestCase):
cache_dir: Path
Expand Down Expand Up @@ -262,3 +271,11 @@ def test_push_to_hub(self):

# Delete repo
self._api.delete_repo(repo_id=repo_id)

def test_save_pretrained_do_not_overwrite_config(self):
"""Regression test for https://github.com/huggingface/huggingface_hub/issues/2102."""
model = DummyModelSavingConfig()
model.save_pretrained(self.cache_dir)
# config.json is not overwritten
with open(self.cache_dir / "config.json") as f:
assert json.load(f) == {"custom_config": "custom_config"}
2 changes: 2 additions & 0 deletions tests/test_hub_mixin_pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,7 @@ def test_load_no_config(self):
assert reloaded_with_default.state == "other"
assert reloaded_with_default.config == {"num_classes": 50, "state": "other"}

config_file.unlink() # Remove config file
reloaded_with_default.save_pretrained(self.cache_dir)
assert json.loads(config_file.read_text()) == {"num_classes": 50, "state": "other"}

Expand All @@ -307,6 +308,7 @@ def test_save_with_non_jsonable_config(self):
assert "not_jsonable" not in model.config

# If jsonable value passed by user, it's saved in the config
(self.cache_dir / "config.json").unlink()
new_model = DummyModelNoConfig(not_jsonable=123)
new_model.save_pretrained(self.cache_dir)
assert new_model.config["not_jsonable"] == 123
Expand Down
Loading