Introducing list of tags to Keras model card

src/huggingface_hub/keras_mixin.py

@@ -1,9 +1,11 @@
 import json
 import os
+import warnings
 from pathlib import Path
 from shutil import copytree, rmtree
 from typing import Any, Dict, Optional, Union
 
+import yaml
 from huggingface_hub import ModelHubMixin
 from huggingface_hub.file_download import (
     is_graphviz_available,
@@ -100,7 +102,7 @@ def _create_model_card(
     model,
     repo_dir: Path,
     plot_model: Optional[bool] = True,
-    task_name: Optional[str] = None,
+    metadata: Optional[dict] = None,
 ):
     """
     Creates a model card for the repository.
@@ -109,10 +111,10 @@ def _create_model_card(
     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"
-    if task_name is not None:
-        model_card += f"tags:\n- {task_name}\n"
-    model_card += "library_name: keras\n---\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"
@@ -149,7 +151,7 @@ def save_pretrained_keras(
     config: Optional[Dict[str, Any]] = None,
     include_optimizer: Optional[bool] = False,
     plot_model: Optional[bool] = True,
-    task_name: Optional[str] = None,
+    tags: Optional[Union[list, str]] = None,
     **model_save_kwargs,
 ):
     """
@@ -170,9 +172,9 @@ def save_pretrained_keras(
         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.
-        task_name (`str`, *optional*):
-            Name of the task the model was trained on. Available tasks
-            [here](https://github.com/huggingface/hub-docs/blob/main/js/src/lib/interfaces/Types.ts).
+        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).
@@ -199,7 +201,24 @@ def save_pretrained_keras(
         with open(path, "w") as f:
             json.dump(config, f)
 
-    _create_model_card(model, save_directory, plot_model, task_name)
+    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]
+
+    _create_model_card(model, save_directory, plot_model, metadata)
     tf.keras.models.save_model(
         model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs
     )
@@ -278,7 +297,7 @@ def push_to_hub_keras(
     git_email: Optional[str] = None,
     config: Optional[dict] = None,
     include_optimizer: Optional[bool] = False,
-    task_name: Optional[str] = None,
+    tags: Optional[Union[list, str]] = None,
     plot_model: Optional[bool] = True,
     **model_save_kwargs,
 ):
@@ -328,9 +347,9 @@ def push_to_hub_keras(
             Configuration object to be saved alongside the model weights.
         include_optimizer (`bool`, *optional*, defaults to `False`):
             Whether or not to include optimizer during serialization.
-        task_name (`str`, *optional*):
-            Name of the task the model was trained on. Available tasks
-            [here](https://github.com/huggingface/huggingface_hub/blob/main/js/src/lib/interfaces/Types.ts).
+        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.
@@ -389,8 +408,8 @@ def push_to_hub_keras(
         repo_path_or_name,
         config=config,
         include_optimizer=include_optimizer,
+        tags=tags,
         plot_model=plot_model,
-        task_name=task_name,
         **model_save_kwargs,
     )
 

tests/test_keras_integration.py

@@ -272,6 +272,23 @@ def test_from_pretrained_weights(self):
             .item()
         )
 
+    def test_save_pretrained_task_name_deprecation(self):
+        REPO_NAME = repo_name("save")
+        model = self.model_init()
+        model.build((None, 2))
+
+        with pytest.warns(
+            FutureWarning,
+            match="`task_name` input argument is removed. Pass `tags` instead.",
+        ):
+            save_pretrained_keras(
+                model,
+                f"{WORKING_REPO_DIR}/{REPO_NAME}",
+                tags=["test"],
+                task_name="test",
+                save_traces=True,
+            )
+
     def test_rel_path_from_pretrained(self):
         model = self.model_init()
         model.build((None, 2))
@@ -300,7 +317,7 @@ def test_abs_path_from_pretrained(self):
             f"{WORKING_REPO_DIR}/{REPO_NAME}",
             config={"num": 10, "act": "gelu_fast"},
             plot_model=True,
-            task_name=None,
+            tags=None,
         )
 
         new_model = from_pretrained_keras(f"{WORKING_REPO_DIR}/{REPO_NAME}")