🚧 wip

timm/models/hub.py

@@ -3,11 +3,10 @@ import logging
 import os
 from functools import partial
 from pathlib import Path
-from typing import Optional, Union
+from typing import Union
 
 import torch
-from torch.hub import (HASH_REGEX, download_url_to_file,
-                       load_state_dict_from_url, urlparse)
+from torch.hub import HASH_REGEX, download_url_to_file, urlparse, load_state_dict_from_url
 
 try:
     from torch.hub import get_dir
@@ -17,8 +16,7 @@ except ImportError:
 from timm import __version__
 
 try:
-    from huggingface_hub import (HfApi, HfFolder, Repository, cached_download,
-                                 hf_hub_url)
+    from huggingface_hub import HfApi, HfFolder, Repository, cached_download, hf_hub_url
     cached_download = partial(cached_download, library_name="timm", library_version=__version__)
 except ImportError:
     hf_hub_url = None
@@ -110,107 +108,57 @@ def save_pretrained_for_hf(model, save_directory, **config_kwargs):
 
     config_path = save_directory / 'config.json'
     config = model.default_cfg
+    config['num_classes'] = config_kwargs.pop('num_classes', model.num_classes)
+    config['num_features'] = config_kwargs.pop('num_features', model.num_features)
+    config['labels'] = config_kwargs.pop('labels', [f"LABEL_{i}" for i in range(config['num_classes'])])
     config.update(config_kwargs)
 
     with config_path.open('w') as f:
-        json.dump(config, f, indent=4)
+        json.dump(config, f, indent=2)
 
 
 def push_to_hf_hub(
     model,
-    repo_path_or_name: Optional[str] = None,
-    repo_url: Optional[str] = None,
-    commit_message: Optional[str] = "Add model",
-    organization: Optional[str] = None,
-    private: Optional[bool] = None,
-    api_endpoint: Optional[str] = None,
-    use_auth_token: Optional[Union[bool, str]] = None,
-    git_user: Optional[str] = None,
-    git_email: Optional[str] = None,
-    config: Optional[dict] = None,
+    local_dir,
+    repo_namespace_or_url=None,
+    commit_message='Add model',
+    use_auth_token=True,
+    git_email=None,
+    git_user=None,
+    revision=None,
+    **config_kwargs
 ):
-    """
-    Upload model checkpoint and config to the 🤗 Model Hub while synchronizing a local clone of the repo in
-    :obj:`repo_path_or_name`.
 
-    Parameters:
-        repo_path_or_name (:obj:`str`, `optional`):
-            Can either be a repository name for your model or tokenizer in the Hub or a path to a local folder (in
-            which case the repository will have the name of that local folder). If not specified, will default to
-            the name given by :obj:`repo_url` and a local directory with that name will be created.
-        repo_url (:obj:`str`, `optional`):
-            Specify this in case you want to push to an existing repository in the hub. If unspecified, a new
-            repository will be created in your namespace (unless you specify an :obj:`organization`) with
-            :obj:`repo_name`.
-        commit_message (:obj:`str`, `optional`):
-            Message to commit while pushing. Will default to :obj:`"add config"`, :obj:`"add tokenizer"` or
-            :obj:`"add model"` depending on the type of the class.
-        organization (:obj:`str`, `optional`):
-            Organization in which you want to push your model or tokenizer (you must be a member of this
-            organization).
-        private (:obj:`bool`, `optional`):
-            Whether or not the repository created should be private (requires a paying subscription).
-        api_endpoint (:obj:`str`, `optional`):
-            The API endpoint to use when pushing the model to the hub.
-        use_auth_token (:obj:`bool` or :obj:`str`, `optional`):
-            The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token
-            generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`). Will default to
-            :obj:`True` if :obj:`repo_url` is not specified.
-        git_user (``str``, `optional`):
-            will override the ``git config user.name`` for committing and pushing files to the hub.
-        git_email (``str``, `optional`):
-            will override the ``git config user.email`` for committing and pushing files to the hub.
-        config (:obj:`dict`, `optional`):
-            Configuration object to be saved alongside the model weights.
+    if repo_namespace_or_url:
+        repo_owner, repo_name = repo_namespace_or_url.rstrip('/').split('/')[-2:]
+    else:
+        if isinstance(use_auth_token, str):
+            token = use_auth_token
+        else:
+            token = HfFolder.get_token()
 
-
-    Returns:
-        The url of the commit of your model in the given repository.
-    """
-    assert has_hf_hub(True)
-    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 use_auth_token is None and repo_url is None:
-        token = HfFolder.get_token()
         if token is None:
             raise ValueError(
                 "You must login to the Hugging Face hub on this computer by typing `transformers-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."
             )
-    elif isinstance(use_auth_token, str):
-        token = use_auth_token
-    else:
-        token = None
 
-    if repo_path_or_name is None:
-        repo_path_or_name = repo_url.split("/")[-1]
+        repo_owner = HfApi().whoami(token)['name']
+        repo_name = Path(local_dir).name
 
-    # 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_name = Path(repo_path_or_name).name
-        repo_url = HfApi(endpoint=api_endpoint).create_repo(
-            token,
-            repo_name,
-            organization=organization,
-            private=private,
-            repo_type=None,
-            exist_ok=True,
-        )
+    repo_url = f'https://huggingface.co/{repo_owner}/{repo_name}'
 
     repo = Repository(
-        repo_path_or_name,
+        local_dir,
         clone_from=repo_url,
         use_auth_token=use_auth_token,
         git_user=git_user,
         git_email=git_email,
+        revision=revision,
     )
-    repo.git_pull(rebase=True)
 
-    save_config = model.default_cfg
-    save_config.update(config or {})
     with repo.commit(commit_message):
-        save_pretrained_for_hf(model, repo.local_dir, **save_config)
+        save_pretrained_for_hf(model, repo.local_dir, **config_kwargs)
+
+    return repo.git_remote_url()