[core/ FEAT] Add the possibility to push custom tags using PreTrainedModel itself

src/transformers/modeling_utils.py

@@ -87,7 +87,7 @@
     replace_return_docstrings,
     strtobool,
 )
-from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files
+from .utils.hub import convert_file_size_to_int, create_and_tag_model_card, get_checkpoint_shard_files
 from .utils.import_utils import (
     ENV_VARS_TRUE_VALUES,
     is_sagemaker_mp_enabled,
@@ -1159,6 +1159,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
     config_class = None
     base_model_prefix = ""
     main_input_name = "input_ids"
+    model_tags = None
+
     _auto_class = None
     _no_split_modules = None
     _skip_keys_device_placement = None
@@ -1239,6 +1241,62 @@ def _backward_compatibility_gradient_checkpointing(self):
             # Remove the attribute now that is has been consumed, so it's no saved in the config.
             delattr(self.config, "gradient_checkpointing")
 
+    def add_model_tags(self, tags: Union[List[str], str]) -> None:
+        r"""
+        Add all tags in `tags` to `model_tags`.
+
+        Args:
+            tags (`Union[List[str], str]`):
+                The desired tags to inject in the model
+        """
+        if isinstance(tags, str):
+            tags = [tags]
+
+        if self.model_tags is None:
+            self.model_tags = []
+
+        for tag in tags:
+            if tag not in self.model_tags:
+                self.model_tags.append(tag)
+
+    def set_model_tags(self, tags: Union[List[str], str]) -> None:
+        r"""
+        Manually force-set the model tags with `tags`
+
+        Args:
+            tags (`Union[List[str], str]`):
+                The desired tags to inject in the model
+        """
+        if isinstance(tags, str):
+            tags = [tags]
+
+        self.model_tags = tags
+
+    def reset_model_tags(self) -> None:
+        r"""
+        Manually reset the model tags with an empty list
+        """
+        if self.model_tags is not None:
+            self.model_tags = []
+
+    def remove_model_tags(self, tags: Union[List[str], str]) -> None:
+        r"""
+        Manually remove all elements of `tags` in the model tags
+
+        Args:
+            tags (`Union[List[str], str]`):
+                The desired tags to remove from the model
+        """
+        if isinstance(tags, str):
+            tags = [tags]
+
+        if self.model_tags is None:
+            return
+
+        for tag in tags:
+            if tag in self.model_tags:
+                self.model_tags.remove(tag)
+
     @classmethod
     def _from_config(cls, config, **kwargs):
         """
@@ -2425,6 +2483,12 @@ def save_pretrained(
             )
 
         if push_to_hub:
+            # Eventually create an empty model card
+            model_card = create_and_tag_model_card(repo_id, self.model_tags)
+
+            # Update model card if needed:
+            model_card.save(os.path.join(save_directory, "README.md"))
+
             self._upload_modified_files(
                 save_directory,
                 repo_id,
@@ -2433,6 +2497,17 @@ def save_pretrained(
                 token=token,
             )
 
+    @wraps(PushToHubMixin.push_to_hub)
+    def push_to_hub(self, *args, **kwargs):
+        if "tags" not in kwargs:
+            kwargs["tags"] = self.model_tags
+        elif "tags" in kwargs and self.model_tags is not None:
+            for model_tag in self.model_tags:
+                # merge the tags together
+                if model_tag not in kwargs["tags"]:
+                    kwargs["tags"].append(model_tag)
+        return super().push_to_hub(*args, **kwargs)
+
     def get_memory_footprint(self, return_buffers=True):
         r"""
         Get the memory footprint of a model. This will return the memory footprint of the current model in bytes.

src/transformers/trainer.py

@@ -3684,6 +3684,18 @@ def push_to_hub(self, commit_message: Optional[str] = "End of training", blockin
         if not self.is_world_process_zero():
             return
 
+        # Add additional tags in the case the model has already some tags and users pass
+        # "tags" argument to `push_to_hub` so that trainer automatically handles internal tags
+        # from all models since Trainer does not call `model.push_to_hub`.
+        if "tags" in kwargs and getattr(self.model, "model_tags", None) is not None:
+            # If it is a string, convert it to an array
+            if isinstance(kwargs["tags"], str):
+                kwargs["tags"] = [kwargs["tags"]]
+
+            for model_tag in self.model.model_tags:
+                if model_tag not in kwargs["tags"]:
+                    kwargs["tags"].append(model_tag)
+
         self.create_model_card(model_name=model_name, **kwargs)
 
         # Wait for the current upload to be finished.

src/transformers/utils/hub.py

@@ -33,6 +33,8 @@
 from huggingface_hub import (
     _CACHED_NO_EXIST,
     CommitOperationAdd,
+    ModelCard,
+    ModelCardData,
     constants,
     create_branch,
     create_commit,
@@ -762,6 +764,7 @@ def push_to_hub(
         safe_serialization: bool = True,
         revision: str = None,
         commit_description: str = None,
+        tags: List[str] = None,
         **deprecated_kwargs,
     ) -> str:
         """
@@ -795,6 +798,8 @@ def push_to_hub(
                 Branch to push the uploaded files to.
             commit_description (`str`, *optional*):
                 The description of the commit that will be created
+            tags (`List[str]`, *optional*):
+                List of tags to push on the Hub.
 
         Examples:
 
@@ -855,6 +860,9 @@ def push_to_hub(
             repo_id, private=private, token=token, repo_url=repo_url, organization=organization
         )
 
+        # Create a new empty model card and eventually tag it
+        model_card = create_and_tag_model_card(repo_id, tags)
+
         if use_temp_dir is None:
             use_temp_dir = not os.path.isdir(working_dir)
 
@@ -864,6 +872,9 @@ def push_to_hub(
             # Save all files.
             self.save_pretrained(work_dir, max_shard_size=max_shard_size, safe_serialization=safe_serialization)
 
+            # Update model card if needed:
+            model_card.save(os.path.join(work_dir, "README.md"))
+
             return self._upload_modified_files(
                 work_dir,
                 repo_id,
@@ -1081,6 +1092,27 @@ def extract_info_from_url(url):
     return {"repo": cache_repo, "revision": revision, "filename": filename}
 
 
+def create_and_tag_model_card(repo_id: str, tags: Optional[List[str]] = None):
+    """
+    Creates a dummy model card and tags it.
+    """
+    try:
+        # Check if the model card is present on the remote repo
+        model_card = ModelCard.load(repo_id)
+    except EntryNotFoundError:
+        # Otherwise create a simple model card from template
+        model_description = "This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated."
+        card_data = ModelCardData(tags=[] if tags is None else tags, library_name="transformers")
+        model_card = ModelCard.from_template(card_data, model_description=model_description)
+
+    if tags is not None:
+        for model_tag in tags:
+            if model_tag not in model_card.data.tags:
+                model_card.data.tags.append(model_tag)
+
+    return model_card
+
+
 def clean_files_for(file):
     """
     Remove, if they exist, file, file.json and file.lock