keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev20240915160609__py3-none-any.whl

keras_hub/src/utils/preset_utils.py

@@ -60,6 +60,8 @@ TOKENIZER_ASSET_DIR = "assets/tokenizer"
 # Config file names.
 CONFIG_FILE = "config.json"
 TOKENIZER_CONFIG_FILE = "tokenizer.json"
+AUDIO_CONVERTER_CONFIG_FILE = "audio_converter.json"
+IMAGE_CONVERTER_CONFIG_FILE = "image_converter.json"
 TASK_CONFIG_FILE = "task.json"
 PREPROCESSOR_CONFIG_FILE = "preprocessor.json"
 METADATA_FILE = "metadata.json"
@@ -77,10 +79,10 @@ SAFETENSOR_FILE = "model.safetensors"
 
 # Global state for preset registry.
 BUILTIN_PRESETS = {}
-BUILTIN_PRESETS_FOR_CLASS = collections.defaultdict(dict)
+BUILTIN_PRESETS_FOR_BACKBONE = collections.defaultdict(dict)
 
 
-def register_presets(presets, classes):
+def register_presets(presets, backbone_cls):
     """Register built-in presets for a set of classes.
 
     Note that this is intended only for models and presets shipped in the
@@ -88,18 +90,26 @@ def register_presets(presets, classes):
     """
     for preset in presets:
         BUILTIN_PRESETS[preset] = presets[preset]
-        for cls in classes:
-            BUILTIN_PRESETS_FOR_CLASS[cls][preset] = presets[preset]
+        BUILTIN_PRESETS_FOR_BACKBONE[backbone_cls][preset] = presets[preset]
 
 
-def list_presets(cls):
+def builtin_presets(cls):
     """Find all registered built-in presets for a class."""
-    return dict(BUILTIN_PRESETS_FOR_CLASS[cls])
+    presets = {}
+    if cls in BUILTIN_PRESETS_FOR_BACKBONE:
+        presets.update(BUILTIN_PRESETS_FOR_BACKBONE[cls])
+    backbone_cls = getattr(cls, "backbone_cls", None)
+    if backbone_cls:
+        presets.update(builtin_presets(backbone_cls))
+    for subclass in list_subclasses(cls):
+        presets.update(builtin_presets(subclass))
+    return presets
 
 
 def list_subclasses(cls):
     """Find all registered subclasses of a class."""
-    custom_objects = keras.saving.get_custom_objects().values()
+    # Deduplicate the lists, since we have to register object twice for compat.
+    custom_objects = set(keras.saving.get_custom_objects().values())
     subclasses = []
     for x in custom_objects:
         if inspect.isclass(x) and x != cls and issubclass(x, cls):
@@ -107,6 +117,26 @@ def list_subclasses(cls):
     return subclasses
 
 
+def find_subclass(preset, cls, backbone_cls):
+    """Find a subclass that is compatible with backbone_cls."""
+    subclasses = list_subclasses(cls)
+    subclasses = filter(lambda x: x.backbone_cls == backbone_cls, subclasses)
+    subclasses = list(subclasses)
+    if not subclasses:
+        raise ValueError(
+            f"Unable to find a subclass of {cls.__name__} that is compatible "
+            f"with {backbone_cls.__name__} found in preset '{preset}'."
+        )
+    # If we find multiple subclasses, try to filter to direct subclasses of
+    # the class we are trying to instantiate.
+    if len(subclasses) > 1:
+        directs = list(filter(lambda x: x in cls.__bases__, subclasses))
+        if len(directs) > 1:
+            subclasses = directs
+    # Return the subclass that was registered first (prefer built-in classes).
+    return subclasses[0]
+
+
 def get_file(preset, path):
     """Download a preset file in necessary and return the local path."""
     # TODO: Add tests for FileNotFound exceptions.
@@ -197,7 +227,7 @@ def get_file(preset, path):
             else:
                 raise ValueError(message)
     elif os.path.exists(preset):
-        # Assume a local filepath.
+        # Assume a local filepath.pyth
         local_path = os.path.join(preset, path)
         if not os.path.exists(local_path):
             raise FileNotFoundError(
@@ -272,6 +302,7 @@ def recursive_pop(config, key):
             recursive_pop(value, key)
 
 
+# TODO: refactor saving routines into a PresetSaver class?
 def make_preset_dir(preset):
     os.makedirs(preset, exist_ok=True)
 
@@ -314,19 +345,9 @@ def save_metadata(layer, preset):
         metadata_file.write(json.dumps(metadata, indent=4))
 
 
-def _validate_tokenizer(preset, allow_incomplete=False):
+def _validate_tokenizer(preset):
     if not check_file_exists(preset, TOKENIZER_CONFIG_FILE):
-        if allow_incomplete:
-            logging.warning(
-                f"`{TOKENIZER_CONFIG_FILE}` is missing from the preset directory `{preset}`."
-            )
-            return
-        else:
-            raise FileNotFoundError(
-                f"`{TOKENIZER_CONFIG_FILE}` is missing from the preset directory `{preset}`. "
-                "To upload the model without a tokenizer, "
-                "set `allow_incomplete=True`."
-            )
+        return
     config_path = get_file(preset, TOKENIZER_CONFIG_FILE)
     try:
         with open(config_path, encoding="utf-8") as config_file:
@@ -377,7 +398,7 @@ def _validate_backbone(preset):
         )
 
 
-def get_snake_case(name):
+def to_snake_case(name):
     name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
     return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()
 
@@ -386,7 +407,7 @@ def create_model_card(preset):
     model_card_path = os.path.join(preset, README_FILE)
     markdown_content = ""
 
-    config = load_config(preset, CONFIG_FILE)
+    config = load_json(preset, CONFIG_FILE)
     model_name = (
         config["class_name"].replace("Backbone", "")
         if config["class_name"].endswith("Backbone")
@@ -395,7 +416,7 @@ def create_model_card(preset):
 
     task_type = None
     if check_file_exists(preset, TASK_CONFIG_FILE):
-        task_config = load_config(preset, TASK_CONFIG_FILE)
+        task_config = load_json(preset, TASK_CONFIG_FILE)
         task_type = (
             task_config["class_name"].replace(model_name, "")
             if task_config["class_name"].startswith(model_name)
@@ -407,12 +428,12 @@ def create_model_card(preset):
     markdown_content += "library_name: keras-hub\n"
     if task_type == "CausalLM":
         markdown_content += "pipeline_tag: text-generation\n"
-    elif task_type == "Classifier":
+    elif task_type == "TextClassifier":
         markdown_content += "pipeline_tag: text-classification\n"
     markdown_content += "---\n"
 
     model_link = (
-        f"https://keras.io/api/keras_hub/models/{get_snake_case(model_name)}"
+        f"https://keras.io/api/keras_hub/models/{to_snake_case(model_name)}"
     )
     markdown_content += (
         f"This is a [`{model_name}` model]({model_link}) "
@@ -454,7 +475,6 @@ def delete_model_card(preset):
 def upload_preset(
     uri,
     preset,
-    allow_incomplete=False,
 ):
     """Upload a preset directory to a model hub.
 
@@ -466,9 +486,6 @@ def upload_preset(
              `hf://[<HF_USERNAME>/]<MODEL>` will be uploaded to the Hugging
              Face Hub.
         preset: The path to the local model preset directory.
-        allow_incomplete: If True, allows the upload of presets without
-                          a tokenizer configuration. Otherwise, a tokenizer
-                          is required.
     """
 
     # Check if preset directory exists.
@@ -476,7 +493,7 @@ def upload_preset(
         raise FileNotFoundError(f"The preset directory {preset} doesn't exist.")
 
     _validate_backbone(preset)
-    _validate_tokenizer(preset, allow_incomplete)
+    _validate_tokenizer(preset)
 
     if uri.startswith(KAGGLE_PREFIX):
         if kagglehub is None:
@@ -533,42 +550,14 @@ def upload_preset(
         )
 
 
-def load_config(preset, config_file=CONFIG_FILE):
+def load_json(preset, config_file=CONFIG_FILE):
     config_path = get_file(preset, config_file)
     with open(config_path, encoding="utf-8") as config_file:
         config = json.load(config_file)
     return config
 
 
-def check_format(preset):
-    if check_file_exists(preset, SAFETENSOR_FILE) or check_file_exists(
-        preset, SAFETENSOR_CONFIG_FILE
-    ):
-        # Determine the format by parsing the config file.
-        config = load_config(preset, HF_CONFIG_FILE)
-        if "hf://timm" in preset or "architecture" in config:
-            return "timm"
-        return "transformers"
-
-    if not check_file_exists(preset, METADATA_FILE):
-        raise FileNotFoundError(
-            f"The preset directory `{preset}` doesn't have a file named `{METADATA_FILE}`, "
-            "or you do not have access to it. This file is required to load a Keras model "
-            "preset. Please verify that the model you are trying to load is a Keras model."
-        )
-    metadata = load_config(preset, METADATA_FILE)
-    if "keras_version" not in metadata:
-        raise ValueError(
-            f"`{METADATA_FILE}` in the preset directory `{preset}` doesn't have `keras_version`. "
-            "Please verify that the model you are trying to load is a Keras model."
-        )
-    return "keras"
-
-
-def load_serialized_object(preset, config_file=CONFIG_FILE, **kwargs):
-    kwargs = kwargs or {}
-    config = load_config(preset, config_file)
-
+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)
@@ -578,14 +567,8 @@ def load_serialized_object(preset, config_file=CONFIG_FILE, **kwargs):
     return keras.saving.deserialize_keras_object(config)
 
 
-def check_config_class(
-    preset,
-    config_file=CONFIG_FILE,
-):
+def check_config_class(config):
     """Validate a preset is being loaded on the correct class."""
-    config_path = get_file(preset, config_file)
-    with open(config_path, encoding="utf-8") as config_file:
-        config = json.load(config_file)
     return keras.saving.get_registered_object(config["registered_name"])
 
 
@@ -619,3 +602,173 @@ def set_dtype_in_config(config, dtype=None):
         for k in policy_map_config["policy_map"].keys():
             policy_map_config["policy_map"][k]["config"]["source_name"] = dtype
     return config
+
+
+def get_preset_loader(preset):
+    if not check_file_exists(preset, CONFIG_FILE):
+        raise ValueError(
+            f"Preset {preset} has no {CONFIG_FILE}. Make sure the URI or "
+            "directory you are trying to load is a valid KerasHub preset and "
+            "and that you have permissions to read/download from this location."
+        )
+    # We currently assume all formats we support have a `config.json`, this is
+    # true, for Keras, Transformers, and timm. We infer the on disk format by
+    # inspecting the `config.json` file.
+    config = load_json(preset, CONFIG_FILE)
+    if "registered_name" in config:
+        # If we see registered_name, we assume a serialized Keras object.
+        return KerasPresetLoader(preset, config)
+    elif "model_type" in config:
+        # Avoid circular import.
+        from keras_hub.src.utils.transformers.preset_loader import (
+            TransformersPresetLoader,
+        )
+
+        # If we see model_type, we assume a Transformers style config.
+        return TransformersPresetLoader(preset, config)
+    elif "architecture" in config:
+        # Avoid circular import.
+        from keras_hub.src.utils.timm.preset_loader import TimmPresetLoader
+
+        # If we see "architecture", we assume a timm config. We could make this
+        # more robust later on if we need to.
+        return TimmPresetLoader(preset, config)
+
+    else:
+        contents = json.dumps(config, indent=4)
+        raise ValueError(
+            f"Unrecognized format for {CONFIG_FILE} in {preset}. "
+            "Create a preset with the `save_to_preset` utility on KerasHub "
+            f"models. Contents of {CONFIG_FILE}:\n{contents}"
+        )
+
+
+class PresetLoader:
+    def __init__(self, preset, config):
+        self.config = config
+        self.preset = preset
+
+    def check_backbone_class(self):
+        """Infer the backbone architecture."""
+        raise NotImplementedError
+
+    def load_backbone(self, cls, load_weights, **kwargs):
+        """Load the backbone model from the preset."""
+        raise NotImplementedError
+
+    def load_tokenizer(self, cls, **kwargs):
+        """Load a tokenizer layer from the preset."""
+        raise NotImplementedError
+
+    def load_audio_converter(self, cls, **kwargs):
+        """Load an audio converter layer from the preset."""
+        raise NotImplementedError
+
+    def load_image_converter(self, cls, **kwargs):
+        """Load an image converter layer from the preset."""
+        raise NotImplementedError
+
+    def load_task(self, cls, load_weights, load_task_weights, **kwargs):
+        """Load a task model from the preset.
+
+        By default, we create a task from a backbone and preprocessor with
+        default arguments. This means
+        """
+        if "backbone" not in kwargs:
+            backbone_class = cls.backbone_cls
+            # Forward dtype to backbone.
+            backbone_kwargs = {"dtype": kwargs.pop("dtype", None)}
+            kwargs["backbone"] = self.load_backbone(
+                backbone_class, load_weights, **backbone_kwargs
+            )
+        if "preprocessor" not in kwargs and cls.preprocessor_cls:
+            kwargs["preprocessor"] = self.load_preprocessor(
+                cls.preprocessor_cls,
+            )
+        return cls(**kwargs)
+
+    def load_preprocessor(self, cls, **kwargs):
+        """Load a prepocessor layer from the preset.
+
+        By default, we create a preprocessor from a tokenizer with default
+        arguments. This allow us to support transformers checkpoints by
+        only converting the backbone and tokenizer.
+        """
+        if "tokenizer" not in kwargs and cls.tokenizer_cls:
+            kwargs["tokenizer"] = self.load_tokenizer(cls.tokenizer_cls)
+        if "audio_converter" not in kwargs and cls.audio_converter_cls:
+            kwargs["audio_converter"] = self.load_audio_converter(
+                cls.audio_converter_cls
+            )
+        if "image_converter" not in kwargs and cls.image_converter_cls:
+            kwargs["image_converter"] = self.load_image_converter(
+                cls.image_converter_cls
+            )
+        return cls(**kwargs)
+
+
+class KerasPresetLoader(PresetLoader):
+    def check_backbone_class(self):
+        return check_config_class(self.config)
+
+    def load_backbone(self, cls, load_weights, **kwargs):
+        backbone = 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, **kwargs):
+        tokenizer_config = load_json(self.preset, TOKENIZER_CONFIG_FILE)
+        tokenizer = load_serialized_object(tokenizer_config, **kwargs)
+        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)
+
+    def load_image_converter(self, cls, **kwargs):
+        converter_config = load_json(self.preset, IMAGE_CONVERTER_CONFIG_FILE)
+        return 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
+        # super class loader.
+        if not check_file_exists(self.preset, TASK_CONFIG_FILE):
+            return super().load_task(
+                cls, load_weights, load_task_weights, **kwargs
+            )
+        task_config = load_json(self.preset, TASK_CONFIG_FILE)
+        if not issubclass(check_config_class(task_config), cls):
+            return super().load_task(
+                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 and task.preprocessor.tokenizer:
+            task.preprocessor.tokenizer.load_preset_assets(self.preset)
+        if load_weights:
+            has_task_weights = check_file_exists(self.preset, TASK_WEIGHTS_FILE)
+            if has_task_weights and load_task_weights:
+                jax_memory_cleanup(task)
+                task_weights = get_file(self.preset, TASK_WEIGHTS_FILE)
+                task.load_task_weights(task_weights)
+            else:
+                jax_memory_cleanup(task.backbone)
+            backbone_weights = get_file(self.preset, MODEL_WEIGHTS_FILE)
+            task.backbone.load_weights(backbone_weights)
+        return task
+
+    def load_preprocessor(self, cls, **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, PREPROCESSOR_CONFIG_FILE):
+            return super().load_preprocessor(cls, **kwargs)
+        preprocessor_json = load_json(self.preset, PREPROCESSOR_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.tokenizer.load_preset_assets(self.preset)
+        return preprocessor