tf-keras-nightly 2.19.0.dev2024112610__py3-none-any.whl → 2.19.0.dev2024113010__py3-none-any.whl

tf_keras/__init__.py

@@ -27,4 +27,4 @@ from tf_keras.src.engine.sequential import Sequential
 from tf_keras.src.engine.training import Model
 
 
-__version__ = "2.19.0.dev2024112610"
+__version__ = "2.19.0.dev2024113010"

tf_keras/src/saving/saving_api.py

@@ -15,6 +15,7 @@
 """Public API surface for saving APIs."""
 
 import os
+import tempfile
 import warnings
 import zipfile
 
@@ -33,17 +34,76 @@ except ImportError:
 is_oss = True
 
 
-def _support_gcs_uri(filepath, save_format, is_oss):
-    """Supports GCS URIs through bigstore via a temporary file."""
-    gs_filepath = None
-    if str(filepath).startswith("gs://") and save_format != "tf":
-        gs_filepath = filepath
-        if not is_oss:
-            gs_filepath = filepath.replace("gs://", "/bigstore/")
-        filepath = os.path.join(
-            saving_lib.get_temp_dir(), os.path.basename(gs_filepath)
-        )
-    return gs_filepath, filepath
+class SupportReadFromRemote:
+    """Supports GCS URIs and other remote paths via a temporary file.
+
+    This is used for `.keras` and H5 files on GCS, CNS and CFS. TensorFlow
+    supports remoted saved model out of the box.
+    """
+
+    def __init__(self, filepath):
+        save_format = get_save_format(filepath, save_format=None)
+        if (
+            saving_lib.is_remote_path(filepath)
+            and not tf.io.gfile.isdir(filepath)
+            and save_format != "tf"
+        ):
+            self.temp_directory = tempfile.TemporaryDirectory()
+            gs_filepath = filepath
+            if not is_oss and str(filepath).startswith("gs://"):
+                gs_filepath = filepath.replace("gs://", "/bigstore/")
+            self.local_filepath = os.path.join(
+                self.temp_directory.name, os.path.basename(filepath)
+            )
+            tf.io.gfile.copy(gs_filepath, self.local_filepath, overwrite=True)
+        else:
+            self.temp_directory = None
+            self.local_filepath = filepath
+
+    def __enter__(self):
+        return self.local_filepath
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.temp_directory is not None:
+            self.temp_directory.cleanup()
+
+
+class SupportWriteToRemote:
+    """Supports GCS URIs and other remote paths via a temporary file.
+
+    This is used for `.keras` and H5 files on GCS, CNS and CFS. TensorFlow
+    supports remoted saved model out of the box.
+    """
+
+    def __init__(self, filepath, overwrite=True, save_format=None):
+        save_format = get_save_format(filepath, save_format=save_format)
+        self.overwrite = overwrite
+        if saving_lib.is_remote_path(filepath) and save_format != "tf":
+            self.temp_directory = tempfile.TemporaryDirectory()
+            self.remote_filepath = filepath
+            if not is_oss and str(filepath).startswith("gs://"):
+                self.remote_filepath = self.remote_filepath.replace(
+                    "gs://", "/bigstore/"
+                )
+            self.local_filepath = os.path.join(
+                self.temp_directory.name, os.path.basename(filepath)
+            )
+        else:
+            self.temp_directory = None
+            self.remote_filepath = None
+            self.local_filepath = filepath
+
+    def __enter__(self):
+        return self.local_filepath
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.temp_directory is not None:
+            tf.io.gfile.copy(
+                self.local_filepath,
+                self.remote_filepath,
+                overwrite=self.overwrite,
+            )
+            self.temp_directory.cleanup()
 
 
 @keras_export("keras.saving.save_model", "keras.models.save_model")
@@ -131,46 +191,49 @@ def save_model(model, filepath, overwrite=True, save_format=None, **kwargs):
     when loading the model. See the `custom_objects` argument in
     `tf.keras.saving.load_model`.
     """
-    save_format = get_save_format(filepath, save_format)
-
-    # Supports GCS URIs through bigstore via a temporary file
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-
-    # Deprecation warnings
-    if save_format == "h5":
-        warnings.warn(
-            "You are saving your model as an HDF5 file via `model.save()`. "
-            "This file format is considered legacy. "
-            "We recommend using instead the native TF-Keras format, "
-            "e.g. `model.save('my_model.keras')`.",
-            stacklevel=2,
-        )
+    # Supports remote paths via a temporary file
+    with SupportWriteToRemote(
+        filepath,
+        overwrite=overwrite,
+        save_format=save_format,
+    ) as local_filepath:
+        save_format = get_save_format(filepath, save_format)
+
+        # Deprecation warnings
+        if save_format == "h5":
+            warnings.warn(
+                "You are saving your model as an HDF5 file via `model.save()`. "
+                "This file format is considered legacy. "
+                "We recommend using instead the native TF-Keras format, "
+                "e.g. `model.save('my_model.keras')`.",
+                stacklevel=2,
+            )
 
-    if save_format == "keras":
-        # If file exists and should not be overwritten.
-        try:
-            exists = os.path.exists(filepath)
-        except TypeError:
-            exists = False
-        if exists and not overwrite:
-            proceed = io_utils.ask_to_proceed_with_overwrite(filepath)
-            if not proceed:
-                return
-        if kwargs:
-            raise ValueError(
-                "The following argument(s) are not supported "
-                f"with the native TF-Keras format: {list(kwargs.keys())}"
+        if save_format == "keras":
+            # If file exists and should not be overwritten.
+            try:
+                exists = os.path.exists(local_filepath)
+            except TypeError:
+                exists = False
+            if exists and not overwrite:
+                proceed = io_utils.ask_to_proceed_with_overwrite(local_filepath)
+                if not proceed:
+                    return
+            if kwargs:
+                raise ValueError(
+                    "The following argument(s) are not supported "
+                    f"with the native TF-Keras format: {list(kwargs.keys())}"
+                )
+            saving_lib.save_model(model, local_filepath)
+        else:
+            # Legacy case
+            return legacy_sm_saving_lib.save_model(
+                model,
+                local_filepath,
+                overwrite=overwrite,
+                save_format=save_format,
+                **kwargs,
             )
-        saving_lib.save_model(model, filepath)
-    else:
-        # Legacy case
-        return legacy_sm_saving_lib.save_model(
-            model,
-            filepath,
-            overwrite=overwrite,
-            save_format=save_format,
-            **kwargs,
-        )
 
 
 @keras_export("keras.saving.load_model", "keras.models.load_model")
@@ -217,94 +280,69 @@ def load_model(
     It is recommended that you use layer attributes to
     access specific variables, e.g. `model.get_layer("dense_1").kernel`.
     """
-    # Supports GCS URIs by copying data to temporary file
-    save_format = get_save_format(filepath, save_format=None)
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-    if gs_filepath is not None:
-        tf.io.gfile.copy(gs_filepath, filepath, overwrite=True)
-
-    is_keras_zip = str(filepath).endswith(".keras") and zipfile.is_zipfile(
-        filepath
-    )
-
-    # Support for remote zip files
-    if (
-        saving_lib.is_remote_path(filepath)
-        and not tf.io.gfile.isdir(filepath)
-        and not is_keras_zip
-    ):
-        local_path = os.path.join(
-            saving_lib.get_temp_dir(), os.path.basename(filepath)
-        )
-
-        # Copy from remote to temporary local directory
-        tf.io.gfile.copy(filepath, local_path, overwrite=True)
-
-        # Switch filepath to local zipfile for loading model
-        if zipfile.is_zipfile(local_path):
-            filepath = local_path
-            is_keras_zip = True
-
-    if is_keras_zip:
-        if kwargs:
-            raise ValueError(
-                "The following argument(s) are not supported "
-                f"with the native TF-Keras format: {list(kwargs.keys())}"
+    # Supports remote paths via a temporary file
+    with SupportReadFromRemote(filepath) as local_filepath:
+        if str(local_filepath).endswith(".keras") and zipfile.is_zipfile(
+            local_filepath
+        ):
+            if kwargs:
+                raise ValueError(
+                    "The following argument(s) are not supported "
+                    f"with the native TF-Keras format: {list(kwargs.keys())}"
+                )
+            return saving_lib.load_model(
+                local_filepath,
+                custom_objects=custom_objects,
+                compile=compile,
+                safe_mode=safe_mode,
             )
-        return saving_lib.load_model(
-            filepath,
+
+        # Legacy case.
+        return legacy_sm_saving_lib.load_model(
+            local_filepath,
             custom_objects=custom_objects,
             compile=compile,
-            safe_mode=safe_mode,
+            **kwargs,
         )
 
-    # Legacy case.
-    return legacy_sm_saving_lib.load_model(
-        filepath, custom_objects=custom_objects, compile=compile, **kwargs
-    )
-
 
 def save_weights(model, filepath, overwrite=True, **kwargs):
-    # Supports GCS URIs through bigstore via a temporary file
-    save_format = get_save_format(filepath, save_format=None)
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-
-    if str(filepath).endswith(".weights.h5"):
-        # If file exists and should not be overwritten.
-        try:
-            exists = os.path.exists(filepath)
-        except TypeError:
-            exists = False
-        if exists and not overwrite:
-            proceed = io_utils.ask_to_proceed_with_overwrite(filepath)
-            if not proceed:
-                return
-        saving_lib.save_weights_only(model, filepath)
-    else:
-        legacy_sm_saving_lib.save_weights(
-            model, filepath, overwrite=overwrite, **kwargs
-        )
+    # Supports remote paths via a temporary file
+    with SupportWriteToRemote(filepath, overwrite=overwrite) as local_filepath:
+        if str(local_filepath).endswith(".weights.h5"):
+            # If file exists and should not be overwritten.
+            try:
+                exists = os.path.exists(local_filepath)
+            except TypeError:
+                exists = False
+            if exists and not overwrite:
+                proceed = io_utils.ask_to_proceed_with_overwrite(local_filepath)
+                if not proceed:
+                    return
+            saving_lib.save_weights_only(model, local_filepath)
+        else:
+            legacy_sm_saving_lib.save_weights(
+                model, local_filepath, overwrite=overwrite, **kwargs
+            )
 
 
 def load_weights(model, filepath, skip_mismatch=False, **kwargs):
-    # Supports GCS URIs by copying data to temporary file
-    save_format = get_save_format(filepath, save_format=None)
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-    if gs_filepath is not None:
-        tf.io.gfile.copy(gs_filepath, filepath, overwrite=True)
-
-    if str(filepath).endswith(".keras") and zipfile.is_zipfile(filepath):
-        saving_lib.load_weights_only(
-            model, filepath, skip_mismatch=skip_mismatch
-        )
-    elif str(filepath).endswith(".weights.h5"):
-        saving_lib.load_weights_only(
-            model, filepath, skip_mismatch=skip_mismatch
-        )
-    else:
-        return legacy_sm_saving_lib.load_weights(
-            model, filepath, skip_mismatch=skip_mismatch, **kwargs
-        )
+    # Supports remote paths via a temporary file
+    with SupportReadFromRemote(filepath) as local_filepath:
+        if str(local_filepath).endswith(".keras") and zipfile.is_zipfile(
+            local_filepath
+        ):
+            saving_lib.load_weights_only(
+                model, local_filepath, skip_mismatch=skip_mismatch
+            )
+        elif str(local_filepath).endswith(".weights.h5"):
+            saving_lib.load_weights_only(
+                model, local_filepath, skip_mismatch=skip_mismatch
+            )
+        else:
+            return legacy_sm_saving_lib.load_weights(
+                model, local_filepath, skip_mismatch=skip_mismatch, **kwargs
+            )
 
 
 def get_save_format(filepath, save_format):

tf_keras/src/saving/saving_lib.py

@@ -156,14 +156,8 @@ def save_model(model, filepath, weights_format="h5"):
             "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
         }
     )
-    # TODO(rameshsampath): Need a better logic for local vs remote path
-    if is_remote_path(filepath):
-        # Remote path. Zip to local memory byte io and copy to remote
-        zip_filepath = io.BytesIO()
-    else:
-        zip_filepath = filepath
     try:
-        with zipfile.ZipFile(zip_filepath, "w") as zf:
+        with zipfile.ZipFile(filepath, "w") as zf:
             with zf.open(_METADATA_FILENAME, "w") as f:
                 f.write(metadata_json.encode())
             with zf.open(_CONFIG_FILENAME, "w") as f:
@@ -195,10 +189,6 @@ def save_model(model, filepath, weights_format="h5"):
             )
             weights_store.close()
             asset_store.close()
-
-        if is_remote_path(filepath):
-            with tf.io.gfile.GFile(filepath, "wb") as f:
-                f.write(zip_filepath.getvalue())
     except Exception as e:
         raise e
     finally:

{tf_keras_nightly-2.19.0.dev2024112610.dist-info → tf_keras_nightly-2.19.0.dev2024113010.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf_keras-nightly
-Version: 2.19.0.dev2024112610
+Version: 2.19.0.dev2024113010
 Summary: Deep learning for humans.
 Home-page: https://keras.io/
 Download-URL: https://github.com/keras-team/tf-keras/tags

{tf_keras_nightly-2.19.0.dev2024112610.dist-info → tf_keras_nightly-2.19.0.dev2024113010.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tf_keras/__init__.py,sha256=NBr1Tzh6ZB8znc9LhSep9zgEMM1EA7B7QrTjkZ6AKDM,911
+tf_keras/__init__.py,sha256=TkJSEfpxF7E_c5XsMLzCNEryzHQURhS-XXilfgzLWS0,911
 tf_keras/__internal__/__init__.py,sha256=OHQbeIC0QtRBI7dgXaJaVbH8F00x8dCI-DvEcIfyMsE,671
 tf_keras/__internal__/backend/__init__.py,sha256=LnMs2A6685gDG79fxqmdulIYlVE_3WmXlBTBo9ZWYcw,162
 tf_keras/__internal__/layers/__init__.py,sha256=F5SGMhOTPzm-PR44VrfinURHcVeQPIEdwnZlAkSTB3A,176
@@ -537,8 +537,8 @@ tf_keras/src/preprocessing/text.py,sha256=aomzwE3G2ErwzgL_Dj3ERA_2k7TZclaPMTBTrW
 tf_keras/src/saving/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 tf_keras/src/saving/object_registration.py,sha256=N8aV6eqREYjW2ueQpL3guYHyh5KXuun3DZAlmjfYrTA,7830
 tf_keras/src/saving/pickle_utils.py,sha256=5GtHzwNWVaYfZ-0zn69-zn2yv3R6JUwzHOOamnjP7r0,2605
-tf_keras/src/saving/saving_api.py,sha256=OepUlpp79IjEy5NdXr6pMQoWvNVU2vTFDDzFojXsRhs,13152
-tf_keras/src/saving/saving_lib.py,sha256=4pgC9vv46losKZ42ytUOGRqxF2PrNuw-o3OQx9b-ylo,24673
+tf_keras/src/saving/saving_api.py,sha256=q2_-CsVg81HeHnXDkRqpoBwI3Mpp_gQXjBidr5V_RLo,14838
+tf_keras/src/saving/saving_lib.py,sha256=Rk5rOvxEmCvwUlG3bS0QpOgtURayYmIAZbh2GeTuUOc,24272
 tf_keras/src/saving/serialization_lib.py,sha256=kX4qf_fRp4LySkH9FU37DMd0AXxiUrXKT-VLR3JPl7w,30152
 tf_keras/src/saving/legacy/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 tf_keras/src/saving/legacy/hdf5_format.py,sha256=IqFXHN96fuqKwu_akaqTyf9ISRPavP3Ahjydat948O4,42438
@@ -606,7 +606,7 @@ tf_keras/src/utils/legacy/__init__.py,sha256=EfMmeHYDzwvxNaktPhQbkTdcPSIGCqMhBND
 tf_keras/utils/__init__.py,sha256=b7_d-USe_EmLo02_P99Q1rUCzKBYayPCfiYFStP-0nw,2735
 tf_keras/utils/experimental/__init__.py,sha256=DzGogE2AosjxOVILQBT8PDDcqbWTc0wWnZRobCdpcec,97
 tf_keras/utils/legacy/__init__.py,sha256=7ujlDa5HeSRcth2NdqA0S1P2-VZF1kB3n68jye6Dj-8,189
-tf_keras_nightly-2.19.0.dev2024112610.dist-info/METADATA,sha256=dqzxkrH3z-LM4Qs7XG7X0diDtwCfWs0GznYSgOtyoG0,1637
-tf_keras_nightly-2.19.0.dev2024112610.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-tf_keras_nightly-2.19.0.dev2024112610.dist-info/top_level.txt,sha256=LC8FK7zHDNKxB17C6lGKvrZ_fZZGJsRiBK23SfiDegY,9
-tf_keras_nightly-2.19.0.dev2024112610.dist-info/RECORD,,
+tf_keras_nightly-2.19.0.dev2024113010.dist-info/METADATA,sha256=0iGYOVdmUKVLmSbD9H2U6BUQFryvRzLQ2Y8yBMNPNdQ,1637
+tf_keras_nightly-2.19.0.dev2024113010.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+tf_keras_nightly-2.19.0.dev2024113010.dist-info/top_level.txt,sha256=LC8FK7zHDNKxB17C6lGKvrZ_fZZGJsRiBK23SfiDegY,9
+tf_keras_nightly-2.19.0.dev2024113010.dist-info/RECORD,,