Repository power-up (#132)

* Repository power-up

* Setup update

* Setup update

* Remove token for the test with remotes

* Expose helpers

* Slight reword

* Add types

* Re-iterate

* Apply Julien's suggestion

Co-authored-by: julien-c <[email protected]>

* push_to_hub rework

* Apply Julien's comments

Co-authored-by: julien-c <[email protected]>

Co-authored-by: Sylvain Gugger <[email protected]>
Co-authored-by: julien-c <[email protected]>

Author: 3 people

src/huggingface_hub/hub_mixin.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+from pathlib import Path
 from typing import Dict, Optional, Union
 
 import requests
@@ -200,74 +201,104 @@ def from_pretrained(
 
         return model
 
-    @staticmethod
     def push_to_hub(
-        save_directory: Optional[str],
-        model_id: Optional[str] = None,
+        self,
+        repo_path_or_name: Optional[str] = None,
         repo_url: Optional[str] = None,
-        commit_message: Optional[str] = "add model",
+        commit_message: Optional[str] = "Add model",
         organization: Optional[str] = None,
-        private: bool = None,
-        api_endpoint=None,
-        use_auth_token: Union[bool, str, None] = 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,
     ) -> str:
         """
+        Upload model checkpoint or tokenizer files to the 🤗 Model Hub while synchronizing a local clone of the repo in
+        :obj:`repo_path_or_name`.
+
         Parameters:
-            save_directory (:obj:`Union[str, os.PathLike]`):
-                Directory having model weights & config.
-            model_id (:obj:`str`, `optional`, defaults to :obj:`save_directory`):
-                Repo name in huggingface_hub. If not specified, repo name will be same as `save_directory`
+            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 existing repo in hub.
+                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.
+                Organization in which you want to push your model or tokenizer (you must be a member of this
+                organization).
             private (:obj:`bool`, `optional`):
-                private: Whether the model repo should be private (requires a paid huggingface.co account)
-            commit_message (:obj:`str`, `optional`, defaults to :obj:`add model`):
-                Message to commit while pushing
+                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 (``str`` or ``bool``, `optional`, defaults ``None``):
-                huggingface_token can be extract from ``HfApi().login(username, password)`` and is used to authenticate
-                 against the hub (useful from Google Colab for instance).
-            git_user (``str``, `optional`, defaults ``None``):
+            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`, defaults ``None``):
+            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:
-            url to commit on remote repo.
+            The url of the commit of your model in the given repository.
         """
-        if model_id is None:
-            model_id = save_directory.split("/")[-1]
-
-        # The auth token is necessary to create a repo
-        if isinstance(use_auth_token, str):
-            huggingface_token = use_auth_token
-        elif use_auth_token is None and repo_url is not None:
-            # If the repo url exists, then no need for a token
-            huggingface_token = None
+
+        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:
-            huggingface_token = HfFolder.get_token()
+            token = None
+
+        if repo_path_or_name is None:
+            repo_path_or_name = repo_url.split("/")[-1]
 
-        if repo_url is None:
+        # 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(
-                huggingface_token,
-                model_id,
+                token,
+                repo_name,
                 organization=organization,
                 private=private,
                 repo_type=None,
                 exist_ok=True,
             )
 
         repo = Repository(
-            save_directory,
+            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 the files in the cloned repo
+        self.save_pretrained(repo_path_or_name, config=config)
 
-        return repo.push_to_hub(commit_message=commit_message)
+        # Commit and push!
+        repo.git_add()
+        repo.git_commit(commit_message)
+        return repo.git_push()

src/huggingface_hub/repository.py

@@ -2,6 +2,7 @@
 import os
 import re
 import subprocess
+from pathlib import Path
 from typing import List, Optional, Union
 
 from .hf_api import HfFolder
@@ -11,6 +12,38 @@
 logger = logging.getLogger(__name__)
 
 
+def is_git_repo(folder: Union[str, Path]):
+    """
+    Check if the folder is the root of a git repository
+    """
+    folder_exists = os.path.exists(os.path.join(folder, ".git"))
+    git_branch = subprocess.run(
+        "git branch".split(), cwd=folder, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
+    return folder_exists and git_branch.returncode == 0
+
+
+def is_local_clone(folder: Union[str, Path], remote_url: str):
+    """
+    Check if the folder is the a local clone of the remote_url
+    """
+    if not is_git_repo(folder):
+        return False
+
+    remotes = subprocess.run(
+        "git remote -v".split(),
+        stderr=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        check=True,
+        encoding="utf-8",
+        cwd=folder,
+    ).stdout
+
+    # Remove token for the test with remotes.
+    remote_url = re.sub(r"https://.*@", "https://", remote_url)
+    return remote_url in remotes.split()
+
+
 class Repository:
     """
     Helper class to wrap the git and git-lfs commands.
@@ -60,7 +93,7 @@ def __init__(
         if clone_from is not None:
             self.clone_from(repo_url=clone_from, use_auth_token=use_auth_token)
         else:
-            if os.path.isdir(os.path.join(self.local_dir, ".git")):
+            if is_git_repo(self.local_dir):
                 logger.debug("[Repository] is a valid git repo")
             else:
                 logger.error(
@@ -109,7 +142,9 @@ def check_git_versions(self):
 
     def clone_from(self, repo_url: str, use_auth_token: Union[bool, str, None] = None):
         """
-        Clone from a remote.
+        Clone from a remote. If the folder already exists, will try to clone the repository within it.
+
+        If this folder is a git repository with linked history, will try to update the repository.
         """
         if isinstance(use_auth_token, str):
             huggingface_token = use_auth_token
@@ -139,6 +174,7 @@ def clone_from(self, repo_url: str, use_auth_token: Union[bool, str, None] = Non
 
             # checks if repository is initialized in a empty repository or in one with files
             if len(os.listdir(self.local_dir)) == 0:
+                logger.debug(f"Cloning {repo_url} into local empty directory.")
                 subprocess.run(
                     ["git", "clone", repo_url, "."],
                     stderr=subprocess.PIPE,
@@ -148,71 +184,39 @@ def clone_from(self, repo_url: str, use_auth_token: Union[bool, str, None] = Non
                     cwd=self.local_dir,
                 )
             else:
-                logger.warning(
-                    "[Repository] local_dir is not empty, so let's try to pull the remote over a non-empty folder."
-                )
-                subprocess.run(
-                    "git init".split(),
-                    stderr=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    check=True,
-                    encoding="utf-8",
-                    cwd=self.local_dir,
-                )
-
-                output = subprocess.run(
-                    "git remote -v".split(),
-                    stderr=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    check=True,
-                    encoding="utf-8",
-                    cwd=self.local_dir,
-                )
-
-                if "origin" not in output.stdout.split():
-                    subprocess.run(
-                        ["git", "remote", "add", "origin", repo_url],
-                        stderr=subprocess.PIPE,
-                        stdout=subprocess.PIPE,
-                        check=True,
-                        encoding="utf-8",
-                        cwd=self.local_dir,
-                    )
-
-                subprocess.run(
-                    "git fetch".split(),
-                    stderr=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    check=True,
-                    encoding="utf-8",
-                    cwd=self.local_dir,
-                )
-                subprocess.run(
-                    "git reset origin/main".split(),
-                    stderr=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    encoding="utf-8",
-                    check=True,
-                    cwd=self.local_dir,
-                )
-
-                output = subprocess.run(
-                    "git branch".split(),
-                    stderr=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    check=True,
-                    encoding="utf-8",
-                    cwd=self.local_dir,
-                )
-
-                if "main" not in output.stdout.split():
-                    subprocess.run(
-                        "git checkout origin/main -t".split(),
-                        stderr=subprocess.PIPE,
-                        stdout=subprocess.PIPE,
-                        encoding="utf-8",
-                        check=True,
-                        cwd=self.local_dir,
+                # Check if the folder is the root of a git repository
+                in_repository = is_git_repo(self.local_dir)
+
+                if in_repository:
+                    if is_local_clone(self.local_dir, repo_url):
+                        logger.debug(
+                            f"{self.local_dir} is already a clone of {repo_url}. Make sure you pull the latest"
+                            "changes with `repo.git_pull()`."
+                        )
+                    else:
+                        output = subprocess.run(
+                            "git remote get-url origin".split(),
+                            stderr=subprocess.PIPE,
+                            stdout=subprocess.PIPE,
+                            encoding="utf-8",
+                            cwd=self.local_dir,
+                        )
+
+                        error_msg = (
+                            f"Tried to clone {repo_url} in an unrelated git repository.\nIf you believe this is an "
+                            f"error, please add a remote with the following URL: {repo_url}."
+                        )
+                        if output.returncode == 0:
+                            error_msg += f"\nLocal path has its origin defined as: {output.stdout}"
+
+                        raise EnvironmentError(error_msg)
+
+                if not in_repository:
+                    raise EnvironmentError(
+                        "Tried to clone a repository in a non-empty folder that isn't a git repository. If you really "
+                        "want to do this, do it manually:\m"
+                        "git init && git remote add origin && git pull origin main\n"
+                        " or clone repo to a new folder and move your existing files there afterwards."
                     )
 
         except subprocess.CalledProcessError as exc:

tests/test_hubmixin.py

@@ -113,18 +113,13 @@ def test_abs_path_from_pretrained(self):
 
     def test_push_to_hub(self):
         model = DummyModel()
-        model.save_pretrained(
-            f"{WORKING_REPO_DIR}/{REPO_NAME}-PUSH_TO_HUB",
-            config={"num": 7, "act": "gelu_fast"},
-        )
-
         model.push_to_hub(
-            f"{WORKING_REPO_DIR}/{REPO_NAME}-PUSH_TO_HUB",
-            f"{REPO_NAME}-PUSH_TO_HUB",
+            repo_path_or_name=f"{WORKING_REPO_DIR}/{REPO_NAME}-PUSH_TO_HUB",
             api_endpoint=ENDPOINT_STAGING,
             use_auth_token=self._token,
             git_user="ci",
             git_email="[email protected]",
+            config={"num": 7, "act": "gelu_fast"},
         )
 
         model_info = self._api.model_info(