feat: save/load TF Keras model (#89)

* feat: save/load TF Keras model
* tensorflow >=2.0

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Borda

README.md

@@ -123,6 +123,43 @@ trainer.fit(BoringModel(), ckpt_path=checkpoint_path)
 
 </details>
 
+<details>
+  <summary>TensorFlow / Keras</summary>
+
+Save model:
+
+```python
+from tensorflow import keras
+
+from litmodels import upload_model
+
+# Define the model
+model = keras.Sequential(
+    [
+        keras.layers.Dense(10, input_shape=(784,), name="dense_1"),
+        keras.layers.Dense(10, name="dense_2"),
+    ]
+)
+
+# Compile the model
+model.compile(optimizer="adam", loss="categorical_crossentropy")
+
+# Save the model
+upload_model("lightning-ai/jirka/sample-tf-keras-model", model=model)
+```
+
+Load model:
+
+```python
+from litmodels import load_model
+
+model_ = load_model(
+    "lightning-ai/jirka/sample-tf-keras-model", download_dir="./my-model"
+)
+```
+
+</details>
+
 <details>
   <summary>SKLearn</summary>
 

_requirements/test.txt

@@ -6,3 +6,4 @@ pytest-mock
 pytorch-lightning >=2.0
 scikit-learn >=1.0
 huggingface-hub >=0.29.0
+tensorflow >=2.0

examples/demo-tensorflow-keras.py

@@ -0,0 +1,19 @@
+from tensorflow import keras
+
+from litmodels import load_model, upload_model
+
+if __name__ == "__main__":
+    # Define the model
+    model = keras.Sequential([
+        keras.layers.Dense(10, input_shape=(784,), name="dense_1"),
+        keras.layers.Dense(10, name="dense_2"),
+    ])
+
+    # Compile the model
+    model.compile(optimizer="adam", loss="categorical_crossentropy")
+
+    # Save the model
+    upload_model("lightning-ai/jirka/sample-tf-keras-model", model=model)
+
+    # Load the model
+    model_ = load_model("lightning-ai/jirka/sample-tf-keras-model", download_dir="./my-model")

src/litmodels/io/gateway.py

@@ -3,15 +3,14 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from lightning_utilities import module_available
-
 from litmodels.io.cloud import download_model_files, upload_model_files
-from litmodels.io.utils import dump_pickle, load_pickle
+from litmodels.io.utils import _KERAS_AVAILABLE, _PYTORCH_AVAILABLE, dump_pickle, load_pickle
 
-if module_available("torch"):
+if _PYTORCH_AVAILABLE:
     import torch
-else:
-    torch = None
+
+if _KERAS_AVAILABLE:
+    from tensorflow import keras
 
 if TYPE_CHECKING:
     from lightning_sdk.models import UploadedModelInfo
@@ -48,12 +47,15 @@ def upload_model(
     # if LightningModule and isinstance(model, LightningModule):
     #     path = os.path.join(staging_dir, f"{model.__class__.__name__}.ckpt")
     #     model.save_checkpoint(path)
-    elif torch and isinstance(model, torch.jit.ScriptModule):
+    elif _PYTORCH_AVAILABLE and isinstance(model, torch.jit.ScriptModule):
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.ts")
         model.save(path)
-    elif torch and isinstance(model, torch.nn.Module):
+    elif _PYTORCH_AVAILABLE and isinstance(model, torch.nn.Module):
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.pth")
         torch.save(model.state_dict(), path)
+    elif _KERAS_AVAILABLE and isinstance(model, keras.models.Model):
+        path = os.path.join(staging_dir, f"{model.__class__.__name__}.keras")
+        model.save(path)
     else:
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.pkl")
         dump_pickle(model=model, path=path)
@@ -110,8 +112,10 @@ def load_model(name: str, download_dir: str = ".") -> Any:
     if len(download_paths) > 1:
         raise NotImplementedError("Downloaded model with multiple files is not supported yet.")
     model_path = Path(download_dir) / download_paths[0]
-    if model_path.suffix.lower() == ".pkl":
-        return load_pickle(path=model_path)
     if model_path.suffix.lower() == ".ts":
         return torch.jit.load(model_path)
+    if model_path.suffix.lower() == ".keras":
+        return keras.models.load_model(model_path)
+    if model_path.suffix.lower() == ".pkl":
+        return load_pickle(path=model_path)
     raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")

src/litmodels/io/utils.py

@@ -3,11 +3,15 @@
 from typing import Any, Union
 
 from lightning_utilities import module_available
+from lightning_utilities.core.imports import RequirementCache
 
-if module_available("joblib"):
+_JOBLIB_AVAILABLE = module_available("joblib")
+_PYTORCH_AVAILABLE = module_available("torch")
+_TENSORFLOW_AVAILABLE = module_available("tensorflow")
+_KERAS_AVAILABLE = RequirementCache("tensorflow >=2.0.0")
+
+if _JOBLIB_AVAILABLE:
     import joblib
-else:
-    joblib = None
 
 
 def dump_pickle(model: Any, path: Union[str, Path]) -> None:
@@ -17,7 +21,7 @@ def dump_pickle(model: Any, path: Union[str, Path]) -> None:
         model: The model to be pickled.
         path: The path where the model will be saved.
     """
-    if joblib is not None:
+    if _JOBLIB_AVAILABLE:
         joblib.dump(model, filename=path, compress=7)
     else:
         with open(path, "wb") as fp:
@@ -33,7 +37,7 @@ def load_pickle(path: Union[str, Path]) -> Any:
     Returns:
         The unpickled model.
     """
-    if joblib is not None:
+    if _JOBLIB_AVAILABLE:
         return joblib.load(path)
     with open(path, "rb") as fp:
         return pickle.load(fp)

tests/integrations/test_cloud.py

@@ -10,10 +10,11 @@
 from lightning_sdk.lightning_cloud.rest_client import GridRestClient
 from lightning_sdk.utils.resolve import _resolve_teamspace
 
-from litmodels import download_model, upload_model
+from litmodels import download_model, load_model, upload_model
 from litmodels.integrations.duplicate import duplicate_hf_model
 from litmodels.integrations.mixins import PickleRegistryMixin, PyTorchRegistryMixin
 from litmodels.io.cloud import _list_available_teamspaces
+from litmodels.io.utils import _KERAS_AVAILABLE
 from tests.integrations import (
     _SKIP_IF_LIGHTNING_BELLOW_2_5_1,
     _SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1,
@@ -315,3 +316,32 @@ def test_list_available_teamspaces():
     assert len(teams) > 0
     # using sanitized teamspace name
     assert f"{LIT_ORG}/oss-litmodels" in teams
+
+
+@pytest.mark.cloud
+@pytest.mark.skipif(
+    not _KERAS_AVAILABLE,
+    reason="TensorFlow Keras is not supported on Windows for now.",
+)
+def test_save_load_tensorflow_keras(tmp_path):
+    from tensorflow import keras
+
+    # Define the model
+    model = keras.Sequential([
+        keras.layers.Dense(10, input_shape=(784,), name="dense_1"),
+        keras.layers.Dense(10, name="dense_2"),
+    ])
+
+    # Compile the model
+    model.compile(optimizer="adam", loss="categorical_crossentropy")
+
+    # model name with random hash
+    teamspace, org_team, model_name = _prepare_variables("tf-keras")
+    upload_model(f"{org_team}/{model_name}", model=model)
+
+    # Load the model
+    model_ = load_model(f"{org_team}/{model_name}", download_dir=str(tmp_path))
+
+    # validate the model
+    assert isinstance(model_, type(model))
+    _cleanup_model(teamspace, model_name, expected_num_versions=1)

tests/test_io_cloud.py

@@ -11,6 +11,7 @@
 import litmodels
 from litmodels import download_model, load_model, upload_model
 from litmodels.io import upload_model_files
+from litmodels.io.utils import _KERAS_AVAILABLE
 
 
 @pytest.mark.parametrize(
@@ -109,3 +110,31 @@ def test_load_model_torch_jit(mock_download_model, tmp_path):
         name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True
     )
     assert isinstance(model, torch.jit.ScriptModule)
+
+
+@pytest.mark.skipif(not _KERAS_AVAILABLE, reason="TensorFlow/Keras is not available")
+@mock.patch("litmodels.io.cloud.sdk_download_model")
+def test_load_model_tf_keras(mock_download_model, tmp_path):
+    from tensorflow import keras
+
+    # create a dummy model file
+    model_file = tmp_path / "dummy_model.keras"
+    # Define the model
+    model = keras.Sequential([
+        keras.layers.Dense(10, input_shape=(784,), name="dense_1"),
+        keras.layers.Dense(10, name="dense_2"),
+    ])
+    model.compile(optimizer="adam", loss="categorical_crossentropy")
+    model.save(model_file)
+    # prepare mocked SDK download function
+    mock_download_model.return_value = [str(model_file.name)]
+
+    # The lit-logger function is just a wrapper around the SDK function
+    model = load_model(
+        name="org-name/teamspace/model-name",
+        download_dir=str(tmp_path),
+    )
+    mock_download_model.assert_called_once_with(
+        name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True
+    )
+    assert isinstance(model, keras.models.Model)