Propagate safe_mode flag to legacy h5 loading code.

Also:
- removed no-op renaming code in legacy saving
- uncommented unit tests in `serialization_lib_test.py`

Author: hertschuh

keras/src/legacy/saving/legacy_h5_format.py

@@ -11,6 +11,7 @@
 from keras.src.legacy.saving import saving_options
 from keras.src.legacy.saving import saving_utils
 from keras.src.saving import object_registration
+from keras.src.saving import serialization_lib
 from keras.src.utils import io_utils
 
 try:
@@ -72,7 +73,9 @@ def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True):
             f.close()
 
 
-def load_model_from_hdf5(filepath, custom_objects=None, compile=True):
+def load_model_from_hdf5(
+    filepath, custom_objects=None, compile=True, safe_mode=True
+):
     """Loads a model saved via `save_model_to_hdf5`.
 
     Args:
@@ -128,7 +131,9 @@ def load_model_from_hdf5(filepath, custom_objects=None, compile=True):
             model_config = model_config.decode("utf-8")
         model_config = json_utils.decode(model_config)
 
-        with saving_options.keras_option_scope(use_legacy_config=True):
+        legacy_scope = saving_options.keras_option_scope(use_legacy_config=True)
+        safe_mode_scope = serialization_lib.SafeModeScope(safe_mode)
+        with legacy_scope, safe_mode_scope:
             model = saving_utils.model_from_config(
                 model_config, custom_objects=custom_objects
             )

keras/src/legacy/saving/legacy_h5_format_test.py

@@ -158,8 +158,13 @@ def test_saving_lambda(self):
 
         temp_filepath = os.path.join(self.get_temp_dir(), "lambda_model.h5")
         legacy_h5_format.save_model_to_hdf5(model, temp_filepath)
-        loaded = legacy_h5_format.load_model_from_hdf5(temp_filepath)
 
+        with self.assertRaisesRegex(ValueError, "Deserializing it is unsafe"):
+            legacy_h5_format.load_model_from_hdf5(temp_filepath)
+
+        loaded = legacy_h5_format.load_model_from_hdf5(
+            temp_filepath, safe_mode=False
+        )
         self.assertAllClose(mean, loaded.layers[1].arguments["mu"])
         self.assertAllClose(std, loaded.layers[1].arguments["std"])
 
@@ -353,8 +358,13 @@ def test_saving_lambda(self):
 
         temp_filepath = os.path.join(self.get_temp_dir(), "lambda_model.h5")
         tf_keras_model.save(temp_filepath)
-        loaded = legacy_h5_format.load_model_from_hdf5(temp_filepath)
 
+        with self.assertRaisesRegex(ValueError, "Deserializing it is unsafe"):
+            legacy_h5_format.load_model_from_hdf5(temp_filepath)
+
+        loaded = legacy_h5_format.load_model_from_hdf5(
+            temp_filepath, safe_mode=False
+        )
         self.assertAllClose(mean, loaded.layers[1].arguments["mu"])
         self.assertAllClose(std, loaded.layers[1].arguments["std"])
 

keras/src/legacy/saving/saving_utils.py

@@ -1,4 +1,3 @@
-import json
 import threading
 
 from absl import logging
@@ -81,10 +80,6 @@ def model_from_config(config, custom_objects=None):
             function_dict["config"]["closure"] = function_config[2]
             config["config"]["function"] = function_dict
 
-    # TODO(nkovela): Swap find and replace args during Keras 3.0 release
-    # Replace keras refs with keras
-    config = _find_replace_nested_dict(config, "keras.", "keras.")
-
     return serialization.deserialize_keras_object(
         config,
         module_objects=MODULE_OBJECTS.ALL_OBJECTS,
@@ -231,13 +226,6 @@ def _deserialize_metric(metric_config):
     return metrics_module.deserialize(metric_config)
 
 
-def _find_replace_nested_dict(config, find, replace):
-    dict_str = json.dumps(config)
-    dict_str = dict_str.replace(find, replace)
-    config = json.loads(dict_str)
-    return config
-
-
 def _resolve_compile_arguments_compat(obj, obj_config, module):
     """Resolves backwards compatibility issues with training config arguments.
 

keras/src/legacy/saving/serialization.py

@@ -2,7 +2,6 @@
 
 import contextlib
 import inspect
-import json
 import threading
 import weakref
 
@@ -485,12 +484,6 @@ def deserialize(config, custom_objects=None):
             arg_spec = inspect.getfullargspec(cls.from_config)
             custom_objects = custom_objects or {}
 
-            # TODO(nkovela): Swap find and replace args during Keras 3.0 release
-            # Replace keras refs with keras
-            cls_config = _find_replace_nested_dict(
-                cls_config, "keras.", "keras."
-            )
-
             if "custom_objects" in arg_spec.args:
                 deserialized_obj = cls.from_config(
                     cls_config,
@@ -565,10 +558,3 @@ def validate_config(config):
 def is_default(method):
     """Check if a method is decorated with the `default` wrapper."""
     return getattr(method, "_is_default", False)
-
-
-def _find_replace_nested_dict(config, find, replace):
-    dict_str = json.dumps(config)
-    dict_str = dict_str.replace(find, replace)
-    config = json.loads(dict_str)
-    return config

keras/src/saving/saving_api.py

@@ -194,7 +194,10 @@ def load_model(filepath, custom_objects=None, compile=True, safe_mode=True):
         )
     if str(filepath).endswith((".h5", ".hdf5")):
         return legacy_h5_format.load_model_from_hdf5(
-            filepath, custom_objects=custom_objects, compile=compile
+            filepath,
+            custom_objects=custom_objects,
+            compile=compile,
+            safe_mode=safe_mode,
         )
     elif str(filepath).endswith(".keras"):
         raise ValueError(

keras/src/saving/serialization_lib_test.py

@@ -175,31 +175,30 @@ def test_lambda_fn(self):
         _, new_obj, _ = self.roundtrip(obj, safe_mode=False)
         self.assertEqual(obj["activation"](3), new_obj["activation"](3))
 
-    # TODO
-    # def test_lambda_layer(self):
-    #     lmbda = keras.layers.Lambda(lambda x: x**2)
-    #     with self.assertRaisesRegex(ValueError, "arbitrary code execution"):
-    #         self.roundtrip(lmbda, safe_mode=True)
-
-    #     _, new_lmbda, _ = self.roundtrip(lmbda, safe_mode=False)
-    #     x = ops.random.normal((2, 2))
-    #     y1 = lmbda(x)
-    #     y2 = new_lmbda(x)
-    #     self.assertAllClose(y1, y2, atol=1e-5)
-
-    # def test_safe_mode_scope(self):
-    #     lmbda = keras.layers.Lambda(lambda x: x**2)
-    #     with serialization_lib.SafeModeScope(safe_mode=True):
-    #         with self.assertRaisesRegex(
-    #             ValueError, "arbitrary code execution"
-    #         ):
-    #             self.roundtrip(lmbda)
-    #     with serialization_lib.SafeModeScope(safe_mode=False):
-    #         _, new_lmbda, _ = self.roundtrip(lmbda)
-    #     x = ops.random.normal((2, 2))
-    #     y1 = lmbda(x)
-    #     y2 = new_lmbda(x)
-    #     self.assertAllClose(y1, y2, atol=1e-5)
+    def test_lambda_layer(self):
+        lmbda = keras.layers.Lambda(lambda x: x**2)
+        with self.assertRaisesRegex(ValueError, "Deserializing it is unsafe"):
+            self.roundtrip(lmbda, safe_mode=True)
+
+        _, new_lmbda, _ = self.roundtrip(lmbda, safe_mode=False)
+        x = ops.random.normal((2, 2))
+        y1 = lmbda(x)
+        y2 = new_lmbda(x)
+        self.assertAllClose(y1, y2, atol=1e-5)
+
+    def test_safe_mode_scope(self):
+        lmbda = keras.layers.Lambda(lambda x: x**2)
+        with serialization_lib.SafeModeScope(safe_mode=True):
+            with self.assertRaisesRegex(
+                ValueError, "Deserializing it is unsafe"
+            ):
+                self.roundtrip(lmbda)
+        with serialization_lib.SafeModeScope(safe_mode=False):
+            _, new_lmbda, _ = self.roundtrip(lmbda)
+        x = ops.random.normal((2, 2))
+        y1 = lmbda(x)
+        y2 = new_lmbda(x)
+        self.assertAllClose(y1, y2, atol=1e-5)
 
     @pytest.mark.requires_trainable_backend
     def test_dict_inputs_outputs(self):