include_optimizer parameter introduced to push_to_hub_keras()

src/huggingface_hub/keras_mixin.py

@@ -17,7 +17,11 @@
 
 
 def save_pretrained_keras(
-    model, save_directory: str, config: Optional[Dict[str, Any]] = None
+    model,
+    save_directory: str,
+    config: Optional[Dict[str, Any]] = None,
+    include_optimizer: Optional[bool] = False,
+    **model_save_kwargs,
 ):
     """Saves a Keras model to save_directory in SavedModel format. Use this if you're using the Functional or Sequential APIs.
 
@@ -27,6 +31,10 @@ def save_pretrained_keras(
         Specify directory in which you want to save the Keras model.
     config (:obj:`dict`, `optional`):
         Configuration object to be saved alongside the model weights.
+    include_optimizer(:obj:`bool`, `optional`):
+        Whether or not to include optimizer in serialization.
+    model_save_kwargs(:obj:`dict`, `optional`):
+        Arguments other than default arguments that can be passed to tf.keras.models.save_model().
     """
     if is_tf_available():
         import tensorflow as tf
@@ -50,7 +58,9 @@ def save_pretrained_keras(
         with open(path, "w") as f:
             json.dump(config, f)
 
-    tf.keras.models.save_model(model, save_directory)
+    tf.keras.models.save_model(
+        model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs
+    )
 
 
 def from_pretrained_keras(*args, **kwargs):
@@ -69,6 +79,8 @@ def push_to_hub_keras(
     git_user: Optional[str] = None,
     git_email: Optional[str] = None,
     config: Optional[dict] = None,
+    include_optimizer: Optional[bool] = False,
+    **model_save_kwargs,
 ):
     """
     Upload model checkpoint or tokenizer files to the 🤗 Model Hub while synchronizing a local clone of the repo in
@@ -104,6 +116,10 @@ def push_to_hub_keras(
             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.
+        include_optimizer (:obj:`bool`, `optional`):
+            Whether or not to include optimizer during serialization.
+        model_save_kwargs(:obj:`dict`, `optional`):
+            Arguments other than default arguments that can be passed to tf.keras.models.save_model().
 
     Returns:
         The url of the commit of your model in the given repository.
@@ -150,7 +166,13 @@ def push_to_hub_keras(
     )
     repo.git_pull(rebase=True)
 
-    save_pretrained_keras(model, repo_path_or_name, config=config)
+    save_pretrained_keras(
+        model,
+        repo_path_or_name,
+        config=config,
+        include_optimizer=include_optimizer,
+        **model_save_kwargs,
+    )
 
     # Commit and push!
     repo.git_add(auto_lfs_track=True)

tests/test_keras_integration.py

@@ -178,7 +178,43 @@ def test_save_pretrained(self):
 
         model.build((None, 2))
 
-        save_pretrained_keras(model, f"{WORKING_REPO_DIR}/{REPO_NAME}")
+        save_pretrained_keras(
+            model,
+            f"{WORKING_REPO_DIR}/{REPO_NAME}",
+        )
+        files = os.listdir(f"{WORKING_REPO_DIR}/{REPO_NAME}")
+
+        self.assertIn("saved_model.pb", files)
+        self.assertIn("keras_metadata.pb", files)
+        self.assertEqual(len(files), 4)
+
+    def test_save_pretrained_optimizer_weights(self):
+        REPO_NAME = repo_name("save")
+        model = self.model_init()
+
+        model.build((None, 2))
+
+        save_pretrained_keras(
+            model, f"{WORKING_REPO_DIR}/{REPO_NAME}", include_optimizer=True
+        )
+        files = os.listdir(f"{WORKING_REPO_DIR}/{REPO_NAME}")
+
+        self.assertIn("saved_model.pb", files)
+        self.assertIn("keras_metadata.pb", files)
+        self.assertEqual(len(files), 4)
+
+    def test_save_pretrained_kwargs(self):
+        REPO_NAME = repo_name("save")
+        model = self.model_init()
+
+        model.build((None, 2))
+
+        save_pretrained_keras(
+            model,
+            f"{WORKING_REPO_DIR}/{REPO_NAME}",
+            include_optimizer=False,
+            save_traces=False,
+        )
         files = os.listdir(f"{WORKING_REPO_DIR}/{REPO_NAME}")
 
         self.assertIn("saved_model.pb", files)
@@ -250,6 +286,30 @@ def test_push_to_hub(self):
             git_user="ci",
             git_email="ci@dummy.com",
             config={"num": 7, "act": "gelu_fast"},
+            include_optimizer=False,
+        )
+
+        model_info = HfApi(endpoint=ENDPOINT_STAGING).model_info(
+            f"{USER}/{REPO_NAME}",
+        )
+        self.assertEqual(model_info.modelId, f"{USER}/{REPO_NAME}")
+
+        self._api.delete_repo(name=f"{REPO_NAME}", token=self._token)
+
+    def test_push_to_hub_model_kwargs(self):
+        REPO_NAME = repo_name("PUSH_TO_HUB")
+        model = self.model_init()
+        model.build((None, 2))
+        push_to_hub_keras(
+            model,
+            repo_path_or_name=f"{WORKING_REPO_DIR}/{REPO_NAME}",
+            api_endpoint=ENDPOINT_STAGING,
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="ci@dummy.com",
+            config={"num": 7, "act": "gelu_fast"},
+            include_optimizer=True,
+            save_traces=False,
         )
 
         model_info = HfApi(endpoint=ENDPOINT_STAGING).model_info(