keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/utils/preset_utils.py

@@ -7,19 +7,10 @@ import re
 
 import keras
 from absl import logging
-from packaging.version import parse
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.utils.keras_utils import print_msg
 
-try:
-    import tensorflow as tf
-except ImportError:
-    raise ImportError(
-        "To use `keras_hub`, please install Tensorflow: `pip install tensorflow`. "
-        "The TensorFlow package is required for data preprocessing with any backend."
-    )
-
 try:
     import kagglehub
     from kagglehub.exceptions import KaggleApiHTTPError
@@ -172,26 +163,13 @@ def get_file(preset, path):
                 )
             else:
                 raise ValueError(message)
-
-    elif scheme in tf.io.gfile.get_registered_schemes():
-        url = os.path.join(preset, path)
-        subdir = preset.replace("://", "_").replace("-", "_").replace("/", "_")
-        filename = os.path.basename(path)
-        subdir = os.path.join(subdir, os.path.dirname(path))
-        try:
-            return copy_gfile_to_cache(
-                filename,
-                url,
-                cache_subdir=os.path.join("models", subdir),
-            )
-        except (tf.errors.PermissionDeniedError, tf.errors.NotFoundError) as e:
-            raise FileNotFoundError(
-                f"`{path}` doesn't exist in preset directory `{preset}`.",
-            ) from e
+    elif scheme in tf_registered_schemes():
+        return tf_copy_gfile_to_cache(preset, path)
     elif scheme == HF_SCHEME:
         if huggingface_hub is None:
             raise ImportError(
-                f"`from_preset()` requires the `huggingface_hub` package to load from '{preset}'. "
+                "`from_preset()` requires the `huggingface_hub` package to "
+                "load from '{preset}'. "
                 "Please install with `pip install huggingface_hub`."
             )
         hf_handle = preset.removeprefix(HF_SCHEME + "://")
@@ -225,7 +203,8 @@ def get_file(preset, path):
         raise ValueError(
             "Unknown preset identifier. A preset must be a one of:\n"
             "1) a built-in preset identifier like `'bert_base_en'`\n"
-            "2) a Kaggle Models handle like `'kaggle://keras/bert/keras/bert_base_en'`\n"
+            "2) a Kaggle Models handle like "
+            "`'kaggle://keras/bert/keras/bert_base_en'`\n"
             "3) a Hugging Face handle like `'hf://username/bert_base_en'`\n"
             "4) a path to a local preset directory like `'./bert_base_en`\n"
             "Use `print(cls.presets.keys())` to view all built-in presets for "
@@ -234,29 +213,48 @@ def get_file(preset, path):
         )
 
 
-def copy_gfile_to_cache(filename, url, cache_subdir):
+def tf_registered_schemes():
+    try:
+        import tensorflow as tf
+
+        return tf.io.gfile.get_registered_schemes()
+    except ImportError:
+        return []
+
+
+def tf_copy_gfile_to_cache(preset, path):
     """Much of this is adapted from get_file of keras core."""
     if "KERAS_HOME" in os.environ:
-        cachdir_base = os.environ.get("KERAS_HOME")
+        base_dir = os.environ.get("KERAS_HOME")
     else:
-        cachdir_base = os.path.expanduser(os.path.join("~", ".keras"))
-    if not os.access(cachdir_base, os.W_OK):
-        cachdir_base = os.path.join("/tmp", ".keras")
-    cachedir = os.path.join(cachdir_base, cache_subdir)
-    os.makedirs(cachedir, exist_ok=True)
-
-    fpath = os.path.join(cachedir, filename)
-    if not os.path.exists(fpath):
+        base_dir = os.path.expanduser(os.path.join("~", ".keras"))
+    if not os.access(base_dir, os.W_OK):
+        base_dir = os.path.join("/tmp", ".keras")
+
+    url = os.path.join(preset, path)
+    model_dir = preset.replace("://", "_").replace("-", "_").replace("/", "_")
+    local_path = os.path.join(base_dir, "models", model_dir, path)
+
+    if not os.path.exists(local_path):
         print_msg(f"Downloading data from {url}")
         try:
-            tf.io.gfile.copy(url, fpath)
+            import tensorflow as tf
+
+            os.make_dirs(os.path.dirname(local_path), exist_ok=True)
+            tf.io.gfile.copy(url, local_path)
         except Exception as e:
             # gfile.copy will leave an empty file after an error.
             # Work around this bug.
-            os.remove(fpath)
+            os.remove(local_path)
+            if isinstance(
+                e, tf.errors.PermissionDeniedError, tf.errors.NotFoundError
+            ):
+                raise FileNotFoundError(
+                    f"`{path}` doesn't exist in preset directory `{preset}`.",
+                ) from e
             raise e
 
-    return fpath
+    return local_path
 
 
 def check_file_exists(preset, path):
@@ -267,64 +265,6 @@ def check_file_exists(preset, path):
     return True
 
 
-def get_tokenizer(layer):
-    """Get the tokenizer from any KerasHub model or layer."""
-    # Avoid circular import.
-    from keras_hub.src.tokenizers.tokenizer import Tokenizer
-
-    if isinstance(layer, Tokenizer):
-        return layer
-    if hasattr(layer, "tokenizer"):
-        return layer.tokenizer
-    if hasattr(layer, "preprocessor"):
-        return getattr(layer.preprocessor, "tokenizer", None)
-    return None
-
-
-def recursive_pop(config, key):
-    """Remove a key from a nested config object"""
-    config.pop(key, None)
-    for value in config.values():
-        if isinstance(value, dict):
-            recursive_pop(value, key)
-
-
-# TODO: refactor saving routines into a PresetSaver class?
-def make_preset_dir(preset):
-    os.makedirs(preset, exist_ok=True)
-
-
-def save_serialized_object(
-    layer,
-    preset,
-    config_file=CONFIG_FILE,
-    config_to_skip=[],
-):
-    make_preset_dir(preset)
-    config_path = os.path.join(preset, config_file)
-    config = keras.saving.serialize_keras_object(layer)
-    config_to_skip += ["compile_config", "build_config"]
-    for c in config_to_skip:
-        recursive_pop(config, c)
-    with open(config_path, "w") as config_file:
-        config_file.write(json.dumps(config, indent=4))
-
-
-def save_metadata(layer, preset):
-    from keras_hub.src.version_utils import __version__ as keras_hub_version
-
-    keras_version = keras.version() if hasattr(keras, "version") else None
-    metadata = {
-        "keras_version": keras_version,
-        "keras_hub_version": keras_hub_version,
-        "parameter_count": layer.count_params(),
-        "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
-    }
-    metadata_path = os.path.join(preset, METADATA_FILE)
-    with open(metadata_path, "w") as metadata_file:
-        metadata_file.write(json.dumps(metadata, indent=4))
-
-
 def _validate_backbone(preset):
     config_path = os.path.join(preset, CONFIG_FILE)
     if not os.path.exists(config_path):
@@ -400,8 +340,8 @@ def create_model_card(preset):
         markdown_content += f"* **{k}:** {v}\n"
     markdown_content += "\n"
     markdown_content += (
-        "This model card has been generated automatically and should be completed "
-        "by the model author. See [Model Cards documentation]"
+        "This model card has been generated automatically and should be "
+        "completed by the model author. See [Model Cards documentation]"
         "(https://huggingface.co/docs/hub/model-cards) for more information.\n"
     )
 
@@ -446,20 +386,16 @@ def upload_preset(
     if uri.startswith(KAGGLE_PREFIX):
         if kagglehub is None:
             raise ImportError(
-                "Uploading a model to Kaggle Hub requires the `kagglehub` package. "
-                "Please install with `pip install kagglehub`."
-            )
-        if parse(kagglehub.__version__) < parse("0.2.4"):
-            raise ImportError(
-                "Uploading a model to Kaggle Hub requires the `kagglehub` package version `0.2.4` or higher. "
-                "Please upgrade with `pip install --upgrade kagglehub`."
+                "Uploading a model to Kaggle Hub requires the `kagglehub` "
+                "package. Please install with `pip install kagglehub`."
             )
         kaggle_handle = uri.removeprefix(KAGGLE_PREFIX)
         kagglehub.model_upload(kaggle_handle, preset)
     elif uri.startswith(HF_PREFIX):
         if huggingface_hub is None:
             raise ImportError(
-                f"`upload_preset()` requires the `huggingface_hub` package to upload to '{uri}'. "
+                f"`upload_preset()` requires the `huggingface_hub` package "
+                f"to upload to '{uri}'. "
                 "Please install with `pip install huggingface_hub`."
             )
         hf_handle = uri.removeprefix(HF_PREFIX)
@@ -471,14 +407,15 @@ def upload_preset(
             raise ValueError(
                 "Unexpected Hugging Face URI. Hugging Face model handles "
                 "should have the form 'hf://[{org}/]{model}'. For example, "
-                "'hf://username/bert_base_en' or 'hf://bert_case_en' to implicitly"
-                f"upload to your user account. Received: URI={uri}."
+                "'hf://username/bert_base_en' or 'hf://bert_case_en' to "
+                f"implicitly upload to your user account. Received: URI={uri}."
             ) from e
         has_model_card = huggingface_hub.file_exists(
             repo_id=repo_url.repo_id, filename=README_FILE
         )
         if not has_model_card:
-            # Remote repo doesn't have a model card so a basic model card is automatically generated.
+            # Remote repo doesn't have a model card so a basic model card is
+            # automatically generated.
             create_model_card(preset)
         try:
             huggingface_hub.upload_folder(
@@ -486,13 +423,14 @@ def upload_preset(
             )
         finally:
             if not has_model_card:
-                # Clean up the preset directory in case user attempts to upload the
-                # preset directory into Kaggle hub as well.
+                # Clean up the preset directory in case user attempts to upload
+                # the preset directory into Kaggle hub as well.
                 delete_model_card(preset)
     else:
         raise ValueError(
             "Unknown URI. An URI must be a one of:\n"
-            "1) a Kaggle Model handle like `'kaggle://<KAGGLE_USERNAME>/<MODEL>/<FRAMEWORK>/<VARIATION>'`\n"
+            "1) a Kaggle Model handle like "
+            "`'kaggle://<KAGGLE_USERNAME>/<MODEL>/<FRAMEWORK>/<VARIATION>'`\n"
             "2) a Hugging Face handle like `'hf://[<HF_USERNAME>/]<MODEL>'`\n"
             f"Received: uri='{uri}'."
         )
@@ -505,19 +443,11 @@ def load_json(preset, config_file=CONFIG_FILE):
     return config
 
 
-def load_serialized_object(config, **kwargs):
-    # `dtype` in config might be a serialized `DTypePolicy` or `DTypePolicyMap`.
-    # Ensure that `dtype` is properly configured.
-    dtype = kwargs.pop("dtype", None)
-    config = set_dtype_in_config(config, dtype)
-
-    config["config"] = {**config["config"], **kwargs}
-    return keras.saving.deserialize_keras_object(config)
-
-
 def check_config_class(config):
     """Validate a preset is being loaded on the correct class."""
     registered_name = config["registered_name"]
+    if registered_name in ("Functional", "Sequential"):
+        return keras.Model
     cls = keras.saving.get_registered_object(registered_name)
     if cls is None:
         raise ValueError(
@@ -600,6 +530,13 @@ def get_preset_loader(preset):
         )
 
 
+def get_preset_saver(preset):
+    # Unlike loading, we only support one form of saving; Keras serialized
+    # configs and saved weights. We keep the rough API structure as loading
+    # just for simplicity.
+    return KerasPresetSaver(preset)
+
+
 class PresetLoader:
     def __init__(self, preset, config):
         self.config = config
@@ -612,10 +549,8 @@ class PresetLoader:
         backbone_kwargs["dtype"] = kwargs.pop("dtype", None)
 
         # Forward `height` and `width` to backbone when using `TextToImage`.
-        if "height" in kwargs:
-            backbone_kwargs["height"] = kwargs.pop("height", None)
-        if "width" in kwargs:
-            backbone_kwargs["width"] = kwargs.pop("width", None)
+        if "image_shape" in kwargs:
+            backbone_kwargs["image_shape"] = kwargs.pop("image_shape", None)
 
         return backbone_kwargs, kwargs
 
@@ -627,7 +562,7 @@ class PresetLoader:
         """Load the backbone model from the preset."""
         raise NotImplementedError
 
-    def load_tokenizer(self, cls, config_name=TOKENIZER_CONFIG_FILE, **kwargs):
+    def load_tokenizer(self, cls, config_file=TOKENIZER_CONFIG_FILE, **kwargs):
         """Load a tokenizer layer from the preset."""
         raise NotImplementedError
 
@@ -658,7 +593,7 @@ class PresetLoader:
         return cls(**kwargs)
 
     def load_preprocessor(
-        self, cls, config_name=PREPROCESSOR_CONFIG_FILE, **kwargs
+        self, cls, config_file=PREPROCESSOR_CONFIG_FILE, **kwargs
     ):
         """Load a prepocessor layer from the preset.
 
@@ -675,25 +610,26 @@ class KerasPresetLoader(PresetLoader):
         return check_config_class(self.config)
 
     def load_backbone(self, cls, load_weights, **kwargs):
-        backbone = load_serialized_object(self.config, **kwargs)
+        backbone = self._load_serialized_object(self.config, **kwargs)
         if load_weights:
             jax_memory_cleanup(backbone)
             backbone.load_weights(get_file(self.preset, MODEL_WEIGHTS_FILE))
         return backbone
 
-    def load_tokenizer(self, cls, config_name=TOKENIZER_CONFIG_FILE, **kwargs):
-        tokenizer_config = load_json(self.preset, config_name)
-        tokenizer = load_serialized_object(tokenizer_config, **kwargs)
-        tokenizer.load_preset_assets(self.preset)
+    def load_tokenizer(self, cls, config_file=TOKENIZER_CONFIG_FILE, **kwargs):
+        tokenizer_config = load_json(self.preset, config_file)
+        tokenizer = self._load_serialized_object(tokenizer_config, **kwargs)
+        if hasattr(tokenizer, "load_preset_assets"):
+            tokenizer.load_preset_assets(self.preset)
         return tokenizer
 
     def load_audio_converter(self, cls, **kwargs):
         converter_config = load_json(self.preset, AUDIO_CONVERTER_CONFIG_FILE)
-        return load_serialized_object(converter_config, **kwargs)
+        return self._load_serialized_object(converter_config, **kwargs)
 
     def load_image_converter(self, cls, **kwargs):
         converter_config = load_json(self.preset, IMAGE_CONVERTER_CONFIG_FILE)
-        return load_serialized_object(converter_config, **kwargs)
+        return self._load_serialized_object(converter_config, **kwargs)
 
     def load_task(self, cls, load_weights, load_task_weights, **kwargs):
         # If there is no `task.json` or it's for the wrong class delegate to the
@@ -708,8 +644,16 @@ class KerasPresetLoader(PresetLoader):
                 cls, load_weights, load_task_weights, **kwargs
             )
         # We found a `task.json` with a complete config for our class.
-        task = load_serialized_object(task_config, **kwargs)
-        if task.preprocessor:
+        # Forward backbone args.
+        backbone_kwargs, kwargs = self.get_backbone_kwargs(**kwargs)
+        if "backbone" in task_config["config"]:
+            backbone_config = task_config["config"]["backbone"]["config"]
+            backbone_config = {**backbone_config, **backbone_kwargs}
+            task_config["config"]["backbone"]["config"] = backbone_config
+        task = self._load_serialized_object(task_config, **kwargs)
+        if task.preprocessor and hasattr(
+            task.preprocessor, "load_preset_assets"
+        ):
             task.preprocessor.load_preset_assets(self.preset)
         if load_weights:
             has_task_weights = check_file_exists(self.preset, TASK_WEIGHTS_FILE)
@@ -724,16 +668,124 @@ class KerasPresetLoader(PresetLoader):
         return task
 
     def load_preprocessor(
-        self, cls, config_name=PREPROCESSOR_CONFIG_FILE, **kwargs
+        self, cls, config_file=PREPROCESSOR_CONFIG_FILE, **kwargs
     ):
         # If there is no `preprocessing.json` or it's for the wrong class,
         # delegate to the super class loader.
-        if not check_file_exists(self.preset, config_name):
+        if not check_file_exists(self.preset, config_file):
             return super().load_preprocessor(cls, **kwargs)
-        preprocessor_json = load_json(self.preset, config_name)
+        preprocessor_json = load_json(self.preset, config_file)
         if not issubclass(check_config_class(preprocessor_json), cls):
             return super().load_preprocessor(cls, **kwargs)
         # We found a `preprocessing.json` with a complete config for our class.
-        preprocessor = load_serialized_object(preprocessor_json, **kwargs)
-        preprocessor.load_preset_assets(self.preset)
+        preprocessor = self._load_serialized_object(preprocessor_json, **kwargs)
+        if hasattr(preprocessor, "load_preset_assets"):
+            preprocessor.load_preset_assets(self.preset)
         return preprocessor
+
+    def _load_serialized_object(self, config, **kwargs):
+        # `dtype` in config might be a serialized `DTypePolicy` or
+        # `DTypePolicyMap`. Ensure that `dtype` is properly configured.
+        dtype = kwargs.pop("dtype", None)
+        config = set_dtype_in_config(config, dtype)
+
+        config["config"] = {**config["config"], **kwargs}
+        return keras.saving.deserialize_keras_object(config)
+
+
+class KerasPresetSaver:
+    def __init__(self, preset_dir):
+        os.makedirs(preset_dir, exist_ok=True)
+        self.preset_dir = preset_dir
+
+    def save_backbone(self, backbone):
+        self._save_serialized_object(backbone, config_file=CONFIG_FILE)
+        backbone_weight_path = os.path.join(self.preset_dir, MODEL_WEIGHTS_FILE)
+        backbone.save_weights(backbone_weight_path)
+        self._save_metadata(backbone)
+
+    def save_tokenizer(self, tokenizer):
+        config_file = TOKENIZER_CONFIG_FILE
+        if hasattr(tokenizer, "config_file"):
+            config_file = tokenizer.config_file
+        self._save_serialized_object(tokenizer, config_file)
+        # Save assets.
+        subdir = config_file.split(".")[0]
+        asset_dir = os.path.join(self.preset_dir, ASSET_DIR, subdir)
+        os.makedirs(asset_dir, exist_ok=True)
+        tokenizer.save_assets(asset_dir)
+
+    def save_audio_converter(self, converter):
+        self._save_serialized_object(converter, AUDIO_CONVERTER_CONFIG_FILE)
+
+    def save_image_converter(self, converter):
+        self._save_serialized_object(converter, IMAGE_CONVERTER_CONFIG_FILE)
+
+    def save_task(self, task):
+        # Save task specific config and weights.
+        self._save_serialized_object(task, TASK_CONFIG_FILE)
+        if task.has_task_weights():
+            task_weight_path = os.path.join(self.preset_dir, TASK_WEIGHTS_FILE)
+            task.save_task_weights(task_weight_path)
+        # Save backbone.
+        if hasattr(task.backbone, "save_to_preset"):
+            task.backbone.save_to_preset(self.preset_dir)
+        else:
+            # Allow saving a `keras.Model` that is not a backbone subclass.
+            self.save_backbone(task.backbone)
+        # Save preprocessor.
+        if task.preprocessor and hasattr(task.preprocessor, "save_to_preset"):
+            task.preprocessor.save_to_preset(self.preset_dir)
+        else:
+            # Allow saving a `keras.Layer` that is not a preprocessor subclass.
+            self.save_preprocessor(task.preprocessor)
+
+    def save_preprocessor(self, preprocessor):
+        config_file = PREPROCESSOR_CONFIG_FILE
+        if hasattr(preprocessor, "config_file"):
+            config_file = preprocessor.config_file
+        self._save_serialized_object(preprocessor, config_file)
+        for layer in preprocessor._flatten_layers(include_self=False):
+            if hasattr(layer, "save_to_preset"):
+                layer.save_to_preset(self.preset_dir)
+
+    def _recursive_pop(self, config, key):
+        """Remove a key from a nested config object"""
+        config.pop(key, None)
+        for value in config.values():
+            if isinstance(value, dict):
+                self._recursive_pop(value, key)
+
+    def _save_serialized_object(self, layer, config_file):
+        config_path = os.path.join(self.preset_dir, config_file)
+        config = keras.saving.serialize_keras_object(layer)
+        config_to_skip = ["compile_config", "build_config"]
+        for key in config_to_skip:
+            self._recursive_pop(config, key)
+        with open(config_path, "w") as config_file:
+            config_file.write(json.dumps(config, indent=4))
+
+    def _save_metadata(self, layer):
+        from keras_hub.src.models.task import Task
+        from keras_hub.src.version_utils import __version__ as keras_hub_version
+
+        # Find all tasks that are compatible with the backbone.
+        # E.g. for `BertBackbone` we would have `TextClassifier` and `MaskedLM`.
+        # For `ResNetBackbone` we would have `ImageClassifier`.
+        tasks = list_subclasses(Task)
+        tasks = filter(lambda x: x.backbone_cls is type(layer), tasks)
+        tasks = [task.__base__.__name__ for task in tasks]
+        # Keep task list alphabetical.
+        tasks = sorted(tasks)
+
+        keras_version = keras.version() if hasattr(keras, "version") else None
+        metadata = {
+            "keras_version": keras_version,
+            "keras_hub_version": keras_hub_version,
+            "parameter_count": layer.count_params(),
+            "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
+            "tasks": tasks,
+        }
+        metadata_path = os.path.join(self.preset_dir, METADATA_FILE)
+        with open(metadata_path, "w") as metadata_file:
+            metadata_file.write(json.dumps(metadata, indent=4))

keras_hub/src/utils/tensor_utils.py

@@ -293,10 +293,10 @@ def any_equal(inputs, values, padding_mask):
 
     Args:
         inputs: Input tensor.
-        values: List or iterable of tensors shaped like `inputs` or broadcastable
-            by bit operators.
-        padding_mask: Tensor with shape compatible with inputs that will condition
-            output.
+        values: List or iterable of tensors shaped like `inputs` or
+            broadcastable by bit operators.
+        padding_mask: Tensor with shape compatible with inputs that will
+            condition output.
 
     Returns:
         A tensor with `inputs` shape where each position is True if it contains

keras_hub/src/utils/timm/convert_densenet.py

@@ -59,9 +59,11 @@ def convert_weights(backbone, loader, timm_config):
     num_stacks = len(backbone.stackwise_num_repeats)
     for stack_index in range(num_stacks):
         for block_idx in range(backbone.stackwise_num_repeats[stack_index]):
-            keras_name = f"stack{stack_index+1}_block{block_idx+1}"
+            keras_name = f"stack{stack_index + 1}_block{block_idx + 1}"
             hf_name = (
-                f"features.denseblock{stack_index+1}.denselayer{block_idx+1}"
+                "features."
+                f"denseblock{stack_index + 1}"
+                f".denselayer{block_idx + 1}"
             )
             port_batch_normalization(f"{keras_name}_1_bn", f"{hf_name}.norm1")
             port_conv2d(f"{keras_name}_1_conv", f"{hf_name}.conv1")
@@ -69,8 +71,8 @@ def convert_weights(backbone, loader, timm_config):
             port_conv2d(f"{keras_name}_2_conv", f"{hf_name}.conv2")
 
     for stack_index in range(num_stacks - 1):
-        keras_transition_name = f"transition{stack_index+1}"
-        hf_transition_name = f"features.transition{stack_index+1}"
+        keras_transition_name = f"transition{stack_index + 1}"
+        hf_transition_name = f"features.transition{stack_index + 1}"
         port_batch_normalization(
             f"{keras_transition_name}_bn", f"{hf_transition_name}.norm"
         )