import collections.abc as collections
import json
import os
import warnings
from functools import wraps
from pathlib import Path
from shutil import copytree
from typing import Any, Dict, List, Optional, Union
from huggingface_hub import ModelHubMixin, snapshot_download
from huggingface_hub.utils import (
get_tf_version,
is_graphviz_available,
is_pydot_available,
is_tf_available,
yaml_dump,
)
from .constants import CONFIG_NAME
from .hf_api import HfApi
from .utils import SoftTemporaryDirectory, logging, validate_hf_hub_args
from .utils._typing import CallableT
logger = logging.get_logger(__name__)
keras = None
if is_tf_available():
# Depending on which version of TensorFlow is installed, we need to import
# keras from the correct location.
# See https://github.com/tensorflow/tensorflow/releases/tag/v2.16.1.
# Note: saving a keras model only works with Keras<3.0.
try:
import tf_keras as keras # type: ignore
except ImportError:
import tensorflow as tf # type: ignore
keras = tf.keras
def _requires_keras_2_model(fn: CallableT) -> CallableT:
# Wrapper to raise if user tries to save a Keras 3.x model
@wraps(fn)
def _inner(model, *args, **kwargs):
if not hasattr(model, "history"): # hacky way to check if model is Keras 2.x
raise NotImplementedError(
f"Cannot use '{fn.__name__}': Keras 3.x is not supported."
" Please save models manually and upload them using `upload_folder` or `huggingface-cli upload`."
)
return fn(model, *args, **kwargs)
return _inner # type: ignore [return-value]
def _flatten_dict(dictionary, parent_key=""):
"""Flatten a nested dictionary.
Reference: https://stackoverflow.com/a/6027615/10319735
Args:
dictionary (`dict`):
The nested dictionary to be flattened.
parent_key (`str`):
The parent key to be prefixed to the children keys.
Necessary for recursing over the nested dictionary.
Returns:
The flattened dictionary.
"""
items = []
for key, value in dictionary.items():
new_key = f"{parent_key}.{key}" if parent_key else key
if isinstance(value, collections.MutableMapping):
items.extend(
_flatten_dict(
value,
new_key,
).items()
)
else:
items.append((new_key, value))
return dict(items)
def _create_hyperparameter_table(model):
"""Parse hyperparameter dictionary into a markdown table."""
table = None
if model.optimizer is not None:
optimizer_params = model.optimizer.get_config()
# flatten the configuration
optimizer_params = _flatten_dict(optimizer_params)
optimizer_params["training_precision"] = keras.mixed_precision.global_policy().name
table = "| Hyperparameters | Value |\n| :-- | :-- |\n"
for key, value in optimizer_params.items():
table += f"| {key} | {value} |\n"
return table
def _plot_network(model, save_directory):
keras.utils.plot_model(
model,
to_file=f"{save_directory}/model.png",
show_shapes=False,
show_dtype=False,
show_layer_names=True,
rankdir="TB",
expand_nested=False,
dpi=96,
layer_range=None,
)
def _create_model_card(
model,
repo_dir: Path,
plot_model: bool = True,
metadata: Optional[dict] = None,
):
"""
Creates a model card for the repository.
Do not overwrite an existing README.md file.
"""
readme_path = repo_dir / "README.md"
if readme_path.exists():
return
hyperparameters = _create_hyperparameter_table(model)
if plot_model and is_graphviz_available() and is_pydot_available():
_plot_network(model, repo_dir)
if metadata is None:
metadata = {}
metadata["library_name"] = "keras"
model_card: str = "---\n"
model_card += yaml_dump(metadata, default_flow_style=False)
model_card += "---\n"
model_card += "\n## Model description\n\nMore information needed\n"
model_card += "\n## Intended uses & limitations\n\nMore information needed\n"
model_card += "\n## Training and evaluation data\n\nMore information needed\n"
if hyperparameters is not None:
model_card += "\n## Training procedure\n"
model_card += "\n### Training hyperparameters\n"
model_card += "\nThe following hyperparameters were used during training:\n\n"
model_card += hyperparameters
model_card += "\n"
if plot_model and os.path.exists(f"{repo_dir}/model.png"):
model_card += "\n ## Model Plot\n"
model_card += "\n"
model_card += "\nView Model Plot
\n"
path_to_plot = "./model.png"
model_card += f"\n\n"
model_card += "\n "
readme_path.write_text(model_card)
@_requires_keras_2_model
def save_pretrained_keras(
model,
save_directory: Union[str, Path],
config: Optional[Dict[str, Any]] = None,
include_optimizer: bool = False,
plot_model: bool = True,
tags: Optional[Union[list, str]] = None,
**model_save_kwargs,
):
"""
Saves a Keras model to save_directory in SavedModel format. Use this if
you're using the Functional or Sequential APIs.
Args:
model (`Keras.Model`):
The [Keras
model](https://www.tensorflow.org/api_docs/python/tf/keras/Model)
you'd like to save. The model must be compiled and built.
save_directory (`str` or `Path`):
Specify directory in which you want to save the Keras model.
config (`dict`, *optional*):
Configuration object to be saved alongside the model weights.
include_optimizer(`bool`, *optional*, defaults to `False`):
Whether or not to include optimizer in serialization.
plot_model (`bool`, *optional*, defaults to `True`):
Setting this to `True` will plot the model and put it in the model
card. Requires graphviz and pydot to be installed.
tags (Union[`str`,`list`], *optional*):
List of tags that are related to model or string of a single tag. See example tags
[here](https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1).
model_save_kwargs(`dict`, *optional*):
model_save_kwargs will be passed to
[`tf.keras.models.save_model()`](https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model).
"""
if keras is None:
raise ImportError("Called a Tensorflow-specific function but could not import it.")
if not model.built:
raise ValueError("Model should be built before trying to save")
save_directory = Path(save_directory)
save_directory.mkdir(parents=True, exist_ok=True)
# saving config
if config:
if not isinstance(config, dict):
raise RuntimeError(f"Provided config to save_pretrained_keras should be a dict. Got: '{type(config)}'")
with (save_directory / CONFIG_NAME).open("w") as f:
json.dump(config, f)
metadata = {}
if isinstance(tags, list):
metadata["tags"] = tags
elif isinstance(tags, str):
metadata["tags"] = [tags]
task_name = model_save_kwargs.pop("task_name", None)
if task_name is not None:
warnings.warn(
"`task_name` input argument is deprecated. Pass `tags` instead.",
FutureWarning,
)
if "tags" in metadata:
metadata["tags"].append(task_name)
else:
metadata["tags"] = [task_name]
if model.history is not None:
if model.history.history != {}:
path = save_directory / "history.json"
if path.exists():
warnings.warn(
"`history.json` file already exists, it will be overwritten by the history of this version.",
UserWarning,
)
with path.open("w", encoding="utf-8") as f:
json.dump(model.history.history, f, indent=2, sort_keys=True)
_create_model_card(model, save_directory, plot_model, metadata)
keras.models.save_model(model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs)
def from_pretrained_keras(*args, **kwargs) -> "KerasModelHubMixin":
r"""
Instantiate a pretrained Keras model from a pre-trained model from the Hub.
The model is expected to be in `SavedModel` format.
Args:
pretrained_model_name_or_path (`str` or `os.PathLike`):
Can be either:
- A string, the `model id` of a pretrained model hosted inside a
model repo on huggingface.co. Valid model ids can be located
at the root-level, like `bert-base-uncased`, or namespaced
under a user or organization name, like
`dbmdz/bert-base-german-cased`.
- You can add `revision` by appending `@` at the end of model_id
simply like this: `dbmdz/bert-base-german-cased@main` Revision
is the specific model version to use. It can be a branch name,
a tag name, or a commit id, since we use a git-based system
for storing models and other artifacts on huggingface.co, so
`revision` can be any identifier allowed by git.
- A path to a `directory` containing model weights saved using
[`~transformers.PreTrainedModel.save_pretrained`], e.g.,
`./my_model_directory/`.
- `None` if you are both providing the configuration and state
dictionary (resp. with keyword arguments `config` and
`state_dict`).
force_download (`bool`, *optional*, defaults to `False`):
Whether to force the (re-)download of the model weights and
configuration files, overriding the cached versions if they exist.
resume_download (`bool`, *optional*, defaults to `False`):
Whether to delete incompletely received files. Will attempt to
resume the download if such a file exists.
proxies (`Dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, e.g.,
`{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The
proxies are used on each request.
token (`str` or `bool`, *optional*):
The token to use as HTTP bearer authorization for remote files. If
`True`, will use the token generated when running `transformers-cli
login` (stored in `~/.huggingface`).
cache_dir (`Union[str, os.PathLike]`, *optional*):
Path to a directory in which a downloaded pretrained model
configuration should be cached if the standard cache should not be
used.
local_files_only(`bool`, *optional*, defaults to `False`):
Whether to only look at local files (i.e., do not try to download
the model).
model_kwargs (`Dict`, *optional*):
model_kwargs will be passed to the model during initialization
Passing `token=True` is required when you want to use a private
model.
"""
return KerasModelHubMixin.from_pretrained(*args, **kwargs)
@validate_hf_hub_args
@_requires_keras_2_model
def push_to_hub_keras(
model,
repo_id: str,
*,
config: Optional[dict] = None,
commit_message: str = "Push Keras model using huggingface_hub.",
private: bool = False,
api_endpoint: Optional[str] = None,
token: Optional[str] = None,
branch: Optional[str] = None,
create_pr: Optional[bool] = None,
allow_patterns: Optional[Union[List[str], str]] = None,
ignore_patterns: Optional[Union[List[str], str]] = None,
delete_patterns: Optional[Union[List[str], str]] = None,
log_dir: Optional[str] = None,
include_optimizer: bool = False,
tags: Optional[Union[list, str]] = None,
plot_model: bool = True,
**model_save_kwargs,
):
"""
Upload model checkpoint to the Hub.
Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
`delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
details.
Args:
model (`Keras.Model`):
The [Keras model](`https://www.tensorflow.org/api_docs/python/tf/keras/Model`) you'd like to push to the
Hub. The model must be compiled and built.
repo_id (`str`):
ID of the repository to push to (example: `"username/my-model"`).
commit_message (`str`, *optional*, defaults to "Add Keras model"):
Message to commit while pushing.
private (`bool`, *optional*, defaults to `False`):
Whether the repository created should be private.
api_endpoint (`str`, *optional*):
The API endpoint to use when pushing the model to the hub.
token (`str`, *optional*):
The token to use as HTTP bearer authorization for remote files. If
not set, will use the token set when logging in with
`huggingface-cli login` (stored in `~/.huggingface`).
branch (`str`, *optional*):
The git branch on which to push the model. This defaults to
the default branch as specified in your repository, which
defaults to `"main"`.
create_pr (`boolean`, *optional*):
Whether or not to create a Pull Request from `branch` with that commit.
Defaults to `False`.
config (`dict`, *optional*):
Configuration object to be saved alongside the model weights.
allow_patterns (`List[str]` or `str`, *optional*):
If provided, only files matching at least one pattern are pushed.
ignore_patterns (`List[str]` or `str`, *optional*):
If provided, files matching any of the patterns are not pushed.
delete_patterns (`List[str]` or `str`, *optional*):
If provided, remote files matching any of the patterns will be deleted from the repo.
log_dir (`str`, *optional*):
TensorBoard logging directory to be pushed. The Hub automatically
hosts and displays a TensorBoard instance if log files are included
in the repository.
include_optimizer (`bool`, *optional*, defaults to `False`):
Whether or not to include optimizer during serialization.
tags (Union[`list`, `str`], *optional*):
List of tags that are related to model or string of a single tag. See example tags
[here](https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1).
plot_model (`bool`, *optional*, defaults to `True`):
Setting this to `True` will plot the model and put it in the model
card. Requires graphviz and pydot to be installed.
model_save_kwargs(`dict`, *optional*):
model_save_kwargs will be passed to
[`tf.keras.models.save_model()`](https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model).
Returns:
The url of the commit of your model in the given repository.
"""
api = HfApi(endpoint=api_endpoint)
repo_id = api.create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True).repo_id
# Push the files to the repo in a single commit
with SoftTemporaryDirectory() as tmp:
saved_path = Path(tmp) / repo_id
save_pretrained_keras(
model,
saved_path,
config=config,
include_optimizer=include_optimizer,
tags=tags,
plot_model=plot_model,
**model_save_kwargs,
)
# If `log_dir` provided, delete remote logs and upload new ones
if log_dir is not None:
delete_patterns = (
[]
if delete_patterns is None
else (
[delete_patterns] # convert `delete_patterns` to a list
if isinstance(delete_patterns, str)
else delete_patterns
)
)
delete_patterns.append("logs/*")
copytree(log_dir, saved_path / "logs")
return api.upload_folder(
repo_type="model",
repo_id=repo_id,
folder_path=saved_path,
commit_message=commit_message,
token=token,
revision=branch,
create_pr=create_pr,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
delete_patterns=delete_patterns,
)
class KerasModelHubMixin(ModelHubMixin):
"""
Implementation of [`ModelHubMixin`] to provide model Hub upload/download
capabilities to Keras models.
```python
>>> import tensorflow as tf
>>> from huggingface_hub import KerasModelHubMixin
>>> class MyModel(tf.keras.Model, KerasModelHubMixin):
... def __init__(self, **kwargs):
... super().__init__()
... self.config = kwargs.pop("config", None)
... self.dummy_inputs = ...
... self.layer = ...
... def call(self, *args):
... return ...
>>> # Initialize and compile the model as you normally would
>>> model = MyModel()
>>> model.compile(...)
>>> # Build the graph by training it or passing dummy inputs
>>> _ = model(model.dummy_inputs)
>>> # Save model weights to local directory
>>> model.save_pretrained("my-awesome-model")
>>> # Push model weights to the Hub
>>> model.push_to_hub("my-awesome-model")
>>> # Download and initialize weights from the Hub
>>> model = MyModel.from_pretrained("username/super-cool-model")
```
"""
def _save_pretrained(self, save_directory):
save_pretrained_keras(self, save_directory)
@classmethod
def _from_pretrained(
cls,
model_id,
revision,
cache_dir,
force_download,
proxies,
resume_download,
local_files_only,
token,
config: Optional[Dict[str, Any]] = None,
**model_kwargs,
):
"""Here we just call [`from_pretrained_keras`] function so both the mixin and
functional APIs stay in sync.
TODO - Some args above aren't used since we are calling
snapshot_download instead of hf_hub_download.
"""
if keras is None:
raise ImportError("Called a TensorFlow-specific function but could not import it.")
# Root is either a local filepath matching model_id or a cached snapshot
if not os.path.isdir(model_id):
storage_folder = snapshot_download(
repo_id=model_id,
revision=revision,
cache_dir=cache_dir,
library_name="keras",
library_version=get_tf_version(),
)
else:
storage_folder = model_id
# TODO: change this in a future PR. We are not returning a KerasModelHubMixin instance here...
model = keras.models.load_model(storage_folder)
# For now, we add a new attribute, config, to store the config loaded from the hub/a local dir.
model.config = config
return model