TensorBoard for Keras models

src/huggingface_hub/keras_mixin.py

@@ -2,6 +2,7 @@
 import logging
 import os
 from pathlib import Path
+from shutil import copytree
 from typing import Any, Dict, Optional, Union
 
 from huggingface_hub import ModelHubMixin
@@ -71,6 +72,7 @@ def push_to_hub_keras(
     model,
     repo_path_or_name: Optional[str] = None,
     repo_url: Optional[str] = None,
+    log_dir: Optional[str] = None,
     commit_message: Optional[str] = "Add model",
     organization: Optional[str] = None,
     private: Optional[bool] = None,
@@ -97,6 +99,9 @@ def push_to_hub_keras(
             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`.
+        log_dir (:obj:`str`, `optional`):
+            TensorBoard logging directory to be pushed. The Hub automatically hosts
+            and displays a TensorBoard instance if log files are included in the repository.
         commit_message (:obj:`str`, `optional`):
             Message to commit while pushing. Will default to :obj:`"add model"`.
         organization (:obj:`str`, `optional`):
@@ -173,7 +178,8 @@ def push_to_hub_keras(
         include_optimizer=include_optimizer,
         **model_save_kwargs,
     )
-
+    if log_dir is not None:
+        copytree(log_dir, f"{repo_path_or_name}/logs")
     # Commit and push!
     repo.git_add(auto_lfs_track=True)
     repo.git_commit(commit_message)

tests/test_keras_integration.py

@@ -292,6 +292,32 @@ def test_push_to_hub(self):
 
         self._api.delete_repo(name=f"{REPO_NAME}", token=self._token)
 
+    def test_push_to_hub_tensorboard(self):
+        os.makedirs(f"{WORKING_REPO_DIR}/tb_log_dir")
+        with open(f"{WORKING_REPO_DIR}/tb_log_dir/tensorboard.txt", "w") as fp:
+            fp.write("Keras FTW")
+        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}",
+            log_dir=f"{WORKING_REPO_DIR}/tb_log_dir",
+            api_endpoint=ENDPOINT_STAGING,
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="ci@dummy.com",
+        )
+        model_info = HfApi(endpoint=ENDPOINT_STAGING).model_info(
+            f"{USER}/{REPO_NAME}",
+        )
+
+        self.assertTrue(
+            "logs/tensorboard.txt" in [f.rfilename for f in model_info.siblings]
+        )
+
+        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()