keras_mixin.py

import collections.abc as collections
import json
import os
import tempfile
import warnings
from pathlib import Path
from shutil import copytree, rmtree
from typing import Any, Dict, List, Optional, Union
from urllib.parse import quote

from huggingface_hub import CommitOperationDelete, 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, DEFAULT_REVISION
from .hf_api import HfApi, _parse_revision_from_pr_url, _prepare_upload_folder_commit
from .repository import Repository
from .utils import HfFolder, logging, validate_hf_hub_args
from .utils._deprecation import _deprecate_arguments, _deprecate_positional_args


logger = logging.get_logger(__name__)

if is_tf_available():
    import tensorflow as tf


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 childer 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."""
    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"
        ] = tf.keras.mixed_precision.global_policy().name
        table = "| Hyperparameters | Value |\n| :-- | :-- |\n"
        for key, value in optimizer_params.items():
            table += f"| {key} | {value} |\n"
    else:
        table = None
    return table


def _plot_network(model, save_directory):
    tf.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: Optional[bool] = True,
    metadata: Optional[dict] = None,
):
    """
    Creates a model card for the repository.
    """
    hyperparameters = _create_hyperparameter_table(model)
    if plot_model and is_graphviz_available() and is_pydot_available():
        _plot_network(model, repo_dir)
    readme_path = f"{repo_dir}/README.md"
    metadata["library_name"] = "keras"
    model_card = "---\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<details>"
        model_card += "\n<summary>View Model Plot</summary>\n"
        path_to_plot = "./model.png"
        model_card += f"\n![Model Image]({path_to_plot})\n"
        model_card += "\n</details>"

    if os.path.exists(readme_path):
        with open(readme_path, "r", encoding="utf8") as f:
            readme = f.read()
    else:
        readme = model_card
    with open(readme_path, "w", encoding="utf-8") as f:
        f.write(readme)


def save_pretrained_keras(
    model,
    save_directory: str,
    config: Optional[Dict[str, Any]] = None,
    include_optimizer: Optional[bool] = False,
    plot_model: Optional[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`):
            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/blame/main/modelcard.md).
        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 is_tf_available():
        import tensorflow as tf
    else:
        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")

    os.makedirs(save_directory, exist_ok=True)

    # saving config
    if config:
        if not isinstance(config, dict):
            raise RuntimeError(
                "Provided config to save_pretrained_keras should be a dict. Got:"
                f" '{type(config)}'"
            )
        path = os.path.join(save_directory, CONFIG_NAME)
        with open(path, "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 = os.path.join(save_directory, "history.json")
            if os.path.exists(path):
                warnings.warn(
                    "`history.json` file already exists, it will be overwritten by the"
                    " history of this version.",
                    UserWarning,
                )
            with open(path, "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)
    tf.keras.models.save_model(
        model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs
    )


def from_pretrained_keras(*args, **kwargs):
    r"""
    Instantiate a pretrained Keras model from a pre-trained model from the Hub.
    The model is expected to be in SavedModel format.```

    Parameters:
        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.
        use_auth_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

    <Tip>

    Passing `use_auth_token=True` is required when you want to use a private
    model.

    </Tip>
    """
    return KerasModelHubMixin.from_pretrained(*args, **kwargs)


@_deprecate_positional_args(version="0.12")
@_deprecate_arguments(
    version="0.12",
    deprecated_args={
        "repo_path_or_name",
        "repo_url",
        "organization",
        "use_auth_token",
        "git_user",
        "git_email",
    },
)
@validate_hf_hub_args
def push_to_hub_keras(
    # NOTE: deprecated signature that will change in 0.12
    model,
    *,
    repo_path_or_name: Optional[str] = None,
    repo_url: Optional[str] = None,
    log_dir: Optional[str] = None,
    commit_message: Optional[str] = "Add model",
    organization: Optional[str] = None,
    private: bool = False,
    api_endpoint: Optional[str] = None,
    use_auth_token: Optional[Union[bool, str]] = True,
    git_user: Optional[str] = None,
    git_email: Optional[str] = None,
    config: Optional[dict] = None,
    include_optimizer: Optional[bool] = False,
    tags: Optional[Union[list, str]] = None,
    plot_model: Optional[bool] = True,
    # NOTE: New arguments since 0.9
    token: Optional[str] = True,
    repo_id: Optional[str] = None,  # optional only until 0.12
    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,
    **model_save_kwargs,
    # TODO (release 0.12): signature must be the following
    # model,
    # repo_id: str,
    # *,
    # commit_message: Optional[str] = "Add model",
    # private: bool = None,
    # api_endpoint: Optional[str] = None,
    # token: Optional[str] = True,
    # branch: Optional[str] = None,
    # create_pr: Optional[bool] = None,
    # config: Optional[dict] = None,
    # allow_patterns: Optional[Union[List[str], str]] = None,
    # ignore_patterns: Optional[Union[List[str], str]] = None,
    # log_dir: Optional[str] = None,
    # include_optimizer: Optional[bool] = False,
    # tags: Optional[Union[list, str]] = None,
    # plot_model: Optional[bool] = True,
    # **model_save_kwargs,
):
    """
    Upload model checkpoint or tokenizer files to the Hub while synchronizing a
    local clone of the repo in `repo_path_or_name`.

    Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be
    pushed to the hub. See [`upload_folder`] reference for more details.

    Parameters:
        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`):
            Repository name to which push
        commit_message (`str`, *optional*, defaults to "Add message"):
            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.
        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/blame/main/modelcard.md).
        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.
    """
    if repo_id is not None:
        api = HfApi(endpoint=api_endpoint)
        api.create_repo(
            repo_id=repo_id,
            repo_type="model",
            token=token,
            private=private,
            exist_ok=True,
        )

        # Push the files to the repo in a single commit
        with tempfile.TemporaryDirectory() 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 is provided, delete old logs + add new ones
            operations = []
            if log_dir is not None:
                # Delete previous log files from Hub
                operations += [
                    CommitOperationDelete(path_in_repo=file)
                    for file in api.list_repo_files(
                        repo_id=repo_id, use_auth_token=token
                    )
                    if file.startswith("logs/")
                ]

                # Copy new log files
                copytree(log_dir, saved_path / "logs")

            # NOTE: `_prepare_upload_folder_commit` and `create_commit` calls are
            #       duplicate code from `upload_folder`. We are not directly using
            #       `upload_folder` since we want to add delete operations to the
            #       commit as well.
            operations += _prepare_upload_folder_commit(
                saved_path,
                path_in_repo="",
                allow_patterns=allow_patterns,
                ignore_patterns=ignore_patterns,
            )
            commit_info = api.create_commit(
                repo_type="model",
                repo_id=repo_id,
                operations=operations,
                commit_message=commit_message,
                token=token,
                revision=branch,
                create_pr=create_pr,
            )
            revision = branch
            if revision is None:
                revision = (
                    quote(_parse_revision_from_pr_url(commit_info.pr_url), safe="")
                    if commit_info.pr_url is not None
                    else DEFAULT_REVISION
                )
            return f"{api.endpoint}/{repo_id}/tree/{revision}/"

    # Repo id is None means we use the deprecated version using Git
    # TODO: remove code between here and `return repo.git_push()` in release 0.12
    if repo_path_or_name is None and repo_url is None:
        raise ValueError("You need to specify a `repo_path_or_name` or a `repo_url`.")

    if isinstance(use_auth_token, bool) and use_auth_token:
        token = HfFolder.get_token()
    elif isinstance(use_auth_token, str):
        token = use_auth_token
    else:
        token = None

    if token is None:
        raise ValueError(
            "You must login to the Hugging Face hub on this computer by typing"
            " `huggingface-cli login` and entering your credentials to use"
            " `use_auth_token=True`. Alternatively, you can pass your own token as the"
            " `use_auth_token` argument."
        )

    if repo_path_or_name is None:
        repo_path_or_name = repo_url.split("/")[-1]

    # If no URL is passed and there's no path to a directory containing files, create a repo
    if repo_url is None and not os.path.exists(repo_path_or_name):
        repo_id = Path(repo_path_or_name).name
        if organization:
            repo_id = f"{organization}/{repo_id}"
        repo_url = HfApi(endpoint=api_endpoint).create_repo(
            repo_id=repo_id,
            token=token,
            private=private,
            repo_type=None,
            exist_ok=True,
        )

    repo = Repository(
        repo_path_or_name,
        clone_from=repo_url,
        use_auth_token=use_auth_token,
        git_user=git_user,
        git_email=git_email,
    )
    repo.git_pull(rebase=True)

    save_pretrained_keras(
        model,
        repo_path_or_name,
        config=config,
        include_optimizer=include_optimizer,
        tags=tags,
        plot_model=plot_model,
        **model_save_kwargs,
    )

    if log_dir is not None:
        if os.path.exists(f"{repo_path_or_name}/logs"):
            rmtree(f"{repo_path_or_name}/logs")
        copytree(log_dir, f"{repo_path_or_name}/logs")

    # Commit and push!
    repo.git_add(auto_lfs_track=True)
    repo.git_commit(commit_message)
    return repo.git_push()


class KerasModelHubMixin(ModelHubMixin):
    """
    Mixin to provide model Hub upload/download capabilities to Keras models.
    Override this class to obtain the following internal methods:
    - `_from_pretrained`, to load a model from the Hub or from local files.
    - `_save_pretrained`, to save a model in the `SavedModel` format.
    """

    def __init__(self, *args, **kwargs):
        """
        Mix this class with your keras-model class for ease process of saving &
        loading from huggingface-hub.


        ```python
        >>> 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 ...


        >>> # Init 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)
        >>> # You can save your model like this
        >>> model.save_pretrained("local_model_dir/", push_to_hub=False)
        >>> # Or, you can push to a new public model repo like this
        >>> model.push_to_hub(
        ...     "super-cool-model",
        ...     git_user="your-hf-username",
        ...     git_email="you@somesite.com",
        ... )

        >>> # Downloading weights from hf-hub & model will be initialized from those weights
        >>> model = MyModel.from_pretrained("username/mymodel@main")
        ```
        """

    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,
        use_auth_token,
        **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 is_tf_available():
            import tensorflow as tf
        else:
            raise ImportError(
                "Called a Tensorflow-specific function but could not import it."
            )

        # TODO - Figure out what to do about these config values. Config is not going to be needed to load model
        cfg = model_kwargs.pop("config", None)

        # 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

        model = tf.keras.models.load_model(storage_folder, **model_kwargs)

        # For now, we add a new attribute, config, to store the config loaded from the hub/a local dir.
        model.config = cfg

        return model