Add Keras version check.

Author: james77777778

keras_hub/src/utils/preset_utils.py

@@ -12,6 +12,7 @@
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.utils.keras_utils import print_msg
+from keras_hub.src.utils.keras_utils import sharded_weights_available
 
 try:
     import kagglehub
@@ -743,6 +744,12 @@ def _load_backbone_weights(self, backbone):
         if has_single_file_weights:
             filepath = get_file(self.preset, MODEL_WEIGHTS_FILE)
         else:
+            if not sharded_weights_available():
+                raise RuntimeError(
+                    "Sharded weights loading is not supported in the current "
+                    f"Keras version {keras.__version__}. "
+                    "Please update to a newer version."
+                )
             filepath = get_file(self.preset, SHARDED_MODEL_WEIGHTS_CONFIG_FILE)
             sharded_filenames = self._get_sharded_filenames(filepath)
             for sharded_filename in sharded_filenames:
@@ -768,7 +775,7 @@ def save_backbone(self, backbone, max_shard_size=None):
         # If the size of the backbone is larger than `MAX_SHARD_SIZE`, save
         # sharded weights.
         max_shard_size = max_shard_size or MAX_SHARD_SIZE
-        if backbone_size_in_gb > max_shard_size:
+        if sharded_weights_available() and backbone_size_in_gb > max_shard_size:
             backbone_sharded_weights_config_path = os.path.join(
                 self.preset_dir, SHARDED_MODEL_WEIGHTS_CONFIG_FILE
             )

keras_hub/src/utils/preset_utils_test.py

@@ -12,6 +12,7 @@
 from keras_hub.src.models.bert.bert_tokenizer import BertTokenizer
 from keras_hub.src.models.gemma.gemma_backbone import GemmaBackbone
 from keras_hub.src.tests.test_case import TestCase
+from keras_hub.src.utils.keras_utils import sharded_weights_available
 from keras_hub.src.utils.preset_utils import CONFIG_FILE
 from keras_hub.src.utils.preset_utils import get_preset_saver
 from keras_hub.src.utils.preset_utils import upload_preset
@@ -20,6 +21,9 @@
 class PresetUtilsTest(TestCase):
     @pytest.mark.large
     def test_sharded_weights(self):
+        if not sharded_weights_available():
+            self.skipTest("Sharded weights are not available.")
+
         # Gemma2 config.
         init_kwargs = {
             "vocabulary_size": 4096,  # 256128