From abf9d51bc3d3b2ca69f569b9f226a3cdf4ec1d10 Mon Sep 17 00:00:00 2001 From: nateraw Date: Tue, 7 Sep 2021 18:39:26 -0600 Subject: [PATCH] :construction: wip --- timm/models/hub.py | 118 +++++++++++++-------------------------------- 1 file changed, 33 insertions(+), 85 deletions(-) diff --git a/timm/models/hub.py b/timm/models/hub.py index f5b041a6..4cdca7f0 100644 --- a/timm/models/hub.py +++ b/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. - - - 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 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() + 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] - - # 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_owner = HfApi().whoami(token)['name'] + repo_name = Path(local_dir).name + + 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()