keras-hub-nightly 0.16.1.dev202409250340__py3-none-any.whl → 0.16.1.dev202409270338__py3-none-any.whl

keras_hub/src/tokenizers/tokenizer.py

@@ -1,30 +1,16 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 import os
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.preprocessing.preprocessing_layer import (
     PreprocessingLayer,
 )
-from keras_hub.src.utils.preset_utils import TOKENIZER_ASSET_DIR
+from keras_hub.src.utils.preset_utils import ASSET_DIR
 from keras_hub.src.utils.preset_utils import TOKENIZER_CONFIG_FILE
 from keras_hub.src.utils.preset_utils import builtin_presets
 from keras_hub.src.utils.preset_utils import find_subclass
 from keras_hub.src.utils.preset_utils import get_file
 from keras_hub.src.utils.preset_utils import get_preset_loader
 from keras_hub.src.utils.preset_utils import save_serialized_object
-from keras_hub.src.utils.preset_utils import save_tokenizer_assets
 from keras_hub.src.utils.python_utils import classproperty
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 
@@ -80,6 +66,7 @@ class Tokenizer(PreprocessingLayer):
     backbone_cls = None
 
     def __init__(self, *args, **kwargs):
+        self.config_name = kwargs.pop("config_name", TOKENIZER_CONFIG_FILE)
         super().__init__(*args, **kwargs)
         self.file_assets = None
 
@@ -187,18 +174,26 @@ class Tokenizer(PreprocessingLayer):
             token = getattr(self, attr)
             setattr(self, f"{attr}_id", self.token_to_id(token))
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "config_name": self.config_name,
+            }
+        )
+        return config
+
     def save_to_preset(self, preset_dir):
         """Save tokenizer to a preset directory.
 
         Args:
             preset_dir: The path to the local model preset directory.
         """
-        save_serialized_object(
-            self,
-            preset_dir,
-            config_file=TOKENIZER_CONFIG_FILE,
-        )
-        save_tokenizer_assets(self, preset_dir)
+        save_serialized_object(self, preset_dir, config_file=self.config_name)
+        subdir = self.config_name.split(".")[0]
+        asset_dir = os.path.join(preset_dir, ASSET_DIR, subdir)
+        os.makedirs(asset_dir, exist_ok=True)
+        self.save_assets(asset_dir)
 
     @preprocessing_function
     def call(self, inputs, *args, training=None, **kwargs):
@@ -207,11 +202,11 @@ class Tokenizer(PreprocessingLayer):
     def load_preset_assets(self, preset):
         asset_path = None
         for asset in self.file_assets:
-            asset_path = get_file(
-                preset, os.path.join(TOKENIZER_ASSET_DIR, asset)
-            )
-        tokenizer_asset_dir = os.path.dirname(asset_path)
-        self.load_assets(tokenizer_asset_dir)
+            subdir = self.config_name.split(".")[0]
+            preset_path = os.path.join(ASSET_DIR, subdir, asset)
+            asset_path = get_file(preset, preset_path)
+        tokenizer_config_name = os.path.dirname(asset_path)
+        self.load_assets(tokenizer_config_name)
 
     @classproperty
     def presets(cls):
@@ -222,6 +217,7 @@ class Tokenizer(PreprocessingLayer):
     def from_preset(
         cls,
         preset,
+        config_name=TOKENIZER_CONFIG_FILE,
         **kwargs,
     ):
         """Instantiate a `keras_hub.models.Tokenizer` from a model preset.
@@ -267,4 +263,4 @@ class Tokenizer(PreprocessingLayer):
         backbone_cls = loader.check_backbone_class()
         if cls.backbone_cls != backbone_cls:
             cls = find_subclass(preset, cls, backbone_cls)
-        return loader.load_tokenizer(cls, **kwargs)
+        return loader.load_tokenizer(cls, config_name, **kwargs)

keras_hub/src/tokenizers/unicode_codepoint_tokenizer.py

@@ -1,18 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch

keras_hub/src/tokenizers/word_piece_tokenizer.py

@@ -1,17 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import os
 import re
 from typing import Iterable

keras_hub/src/tokenizers/word_piece_tokenizer_trainer.py

@@ -1,18 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.tokenizers.word_piece_tokenizer import pretokenize
 

keras_hub/src/utils/__init__.py

@@ -1,13 +0,0 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

keras_hub/src/utils/imagenet/__init__.py

@@ -1,13 +0,0 @@
-# Copyright 2024 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

keras_hub/src/utils/imagenet/imagenet_utils.py

@@ -1,16 +1,3 @@
-# Copyright 2024 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export

keras_hub/src/utils/keras_utils.py

@@ -1,17 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import sys
 
 import keras

keras_hub/src/utils/pipeline_model.py

@@ -1,17 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import functools
 import math
 

keras_hub/src/utils/preset_utils.py

@@ -1,17 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import collections
 import datetime
 import inspect
@@ -55,7 +41,8 @@ KAGGLE_SCHEME = "kaggle"
 GS_SCHEME = "gs"
 HF_SCHEME = "hf"
 
-TOKENIZER_ASSET_DIR = "assets/tokenizer"
+ASSET_DIR = "assets"
+TOKENIZER_ASSET_DIR = f"{ASSET_DIR}/tokenizer"
 
 # Config file names.
 CONFIG_FILE = "config.json"
@@ -307,13 +294,6 @@ def make_preset_dir(preset):
     os.makedirs(preset, exist_ok=True)
 
 
-def save_tokenizer_assets(tokenizer, preset):
-    if tokenizer:
-        asset_dir = os.path.join(preset, TOKENIZER_ASSET_DIR)
-        os.makedirs(asset_dir, exist_ok=True)
-        tokenizer.save_assets(asset_dir)
-
-
 def save_serialized_object(
     layer,
     preset,
@@ -345,37 +325,6 @@ def save_metadata(layer, preset):
         metadata_file.write(json.dumps(metadata, indent=4))
 
 
-def _validate_tokenizer(preset):
-    if not check_file_exists(preset, TOKENIZER_CONFIG_FILE):
-        return
-    config_path = get_file(preset, TOKENIZER_CONFIG_FILE)
-    try:
-        with open(config_path, encoding="utf-8") as config_file:
-            config = json.load(config_file)
-    except Exception as e:
-        raise ValueError(
-            f"Tokenizer config file `{config_path}` is an invalid json file. "
-            f"Error message: {e}"
-        )
-    layer = keras.saving.deserialize_keras_object(config)
-
-    for asset in layer.file_assets:
-        asset_path = get_file(preset, os.path.join(TOKENIZER_ASSET_DIR, asset))
-        if not os.path.exists(asset_path):
-            tokenizer_asset_dir = os.path.dirname(asset_path)
-            raise FileNotFoundError(
-                f"Asset `{asset}` doesn't exist in the tokenizer asset direcotry"
-                f" `{tokenizer_asset_dir}`."
-            )
-    config_dir = os.path.dirname(config_path)
-    asset_dir = os.path.join(config_dir, TOKENIZER_ASSET_DIR)
-
-    tokenizer = get_tokenizer(layer)
-    if not tokenizer:
-        raise ValueError(f"Model or layer `{layer}` is missing tokenizer.")
-    tokenizer.load_assets(asset_dir)
-
-
 def _validate_backbone(preset):
     config_path = os.path.join(preset, CONFIG_FILE)
     if not os.path.exists(config_path):
@@ -493,7 +442,6 @@ def upload_preset(
         raise FileNotFoundError(f"The preset directory {preset} doesn't exist.")
 
     _validate_backbone(preset)
-    _validate_tokenizer(preset)
 
     if uri.startswith(KAGGLE_PREFIX):
         if kagglehub is None:
@@ -657,6 +605,20 @@ class PresetLoader:
         self.config = config
         self.preset = preset
 
+    def get_backbone_kwargs(self, **kwargs):
+        backbone_kwargs = {}
+
+        # Forward `dtype` to backbone.
+        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)
+
+        return backbone_kwargs, kwargs
+
     def check_backbone_class(self):
         """Infer the backbone architecture."""
         raise NotImplementedError
@@ -665,7 +627,7 @@ class PresetLoader:
         """Load the backbone model from the preset."""
         raise NotImplementedError
 
-    def load_tokenizer(self, cls, **kwargs):
+    def load_tokenizer(self, cls, config_name=TOKENIZER_CONFIG_FILE, **kwargs):
         """Load a tokenizer layer from the preset."""
         raise NotImplementedError
 
@@ -685,8 +647,7 @@ class PresetLoader:
         """
         if "backbone" not in kwargs:
             backbone_class = cls.backbone_cls
-            # Forward dtype to backbone.
-            backbone_kwargs = {"dtype": kwargs.pop("dtype", None)}
+            backbone_kwargs, kwargs = self.get_backbone_kwargs(**kwargs)
             kwargs["backbone"] = self.load_backbone(
                 backbone_class, load_weights, **backbone_kwargs
             )
@@ -696,23 +657,16 @@ class PresetLoader:
             )
         return cls(**kwargs)
 
-    def load_preprocessor(self, cls, **kwargs):
+    def load_preprocessor(
+        self, cls, config_name=PREPROCESSOR_CONFIG_FILE, **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
-            )
+        kwargs = cls._add_missing_kwargs(self, kwargs)
         return cls(**kwargs)
 
 
@@ -727,8 +681,8 @@ class KerasPresetLoader(PresetLoader):
             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)
+    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)
         return tokenizer
@@ -755,8 +709,8 @@ class KerasPresetLoader(PresetLoader):
             )
         # 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 task.preprocessor:
+            task.preprocessor.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:
@@ -769,15 +723,17 @@ class KerasPresetLoader(PresetLoader):
             task.backbone.load_weights(backbone_weights)
         return task
 
-    def load_preprocessor(self, cls, **kwargs):
+    def load_preprocessor(
+        self, cls, config_name=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, PREPROCESSOR_CONFIG_FILE):
+        if not check_file_exists(self.preset, config_name):
             return super().load_preprocessor(cls, **kwargs)
-        preprocessor_json = load_json(self.preset, PREPROCESSOR_CONFIG_FILE)
+        preprocessor_json = load_json(self.preset, config_name)
         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)
+        preprocessor.load_preset_assets(self.preset)
         return preprocessor

keras_hub/src/utils/python_utils.py

@@ -1,16 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 """Utilities with miscellaneous python extensions."""
 
 

keras_hub/src/utils/tensor_utils.py

@@ -1,17 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import contextlib
 import functools
 import inspect

keras_hub/src/utils/timm/__init__.py

@@ -1,13 +0,0 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

keras_hub/src/utils/timm/convert_densenet.py

@@ -0,0 +1,107 @@
+# Copyright 2024 The KerasHUB Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+
+from keras_hub.src.models.densenet.densenet_backbone import DenseNetBackbone
+
+backbone_cls = DenseNetBackbone
+
+
+def convert_backbone_config(timm_config):
+    timm_architecture = timm_config["architecture"]
+
+    if timm_architecture == "densenet121":
+        stackwise_num_repeats = [6, 12, 24, 16]
+    elif timm_architecture == "densenet169":
+        stackwise_num_repeats = [6, 12, 32, 32]
+    elif timm_architecture == "densenet201":
+        stackwise_num_repeats = [6, 12, 48, 32]
+    else:
+        raise ValueError(
+            f"Currently, the architecture {timm_architecture} is not supported."
+        )
+    return dict(
+        stackwise_num_repeats=stackwise_num_repeats,
+        compression_ratio=0.5,
+        growth_rate=32,
+    )
+
+
+def convert_weights(backbone, loader, timm_config):
+    def port_conv2d(keras_layer_name, hf_weight_prefix):
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).kernel,
+            hf_weight_key=f"{hf_weight_prefix}.weight",
+            hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
+        )
+
+    def port_batch_normalization(keras_layer_name, hf_weight_prefix):
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).gamma,
+            hf_weight_key=f"{hf_weight_prefix}.weight",
+        )
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).beta,
+            hf_weight_key=f"{hf_weight_prefix}.bias",
+        )
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).moving_mean,
+            hf_weight_key=f"{hf_weight_prefix}.running_mean",
+        )
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).moving_variance,
+            hf_weight_key=f"{hf_weight_prefix}.running_var",
+        )
+
+    # Stem
+    port_conv2d("conv1_conv", "features.conv0")
+    port_batch_normalization("conv1_bn", "features.norm0")
+
+    # Stages
+    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}"
+            hf_name = (
+                f"features.denseblock{stack_index+1}.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")
+            port_batch_normalization(f"{keras_name}_2_bn", f"{hf_name}.norm2")
+            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}"
+        port_batch_normalization(
+            f"{keras_transition_name}_bn", f"{hf_transition_name}.norm"
+        )
+        port_conv2d(
+            f"{keras_transition_name}_conv", f"{hf_transition_name}.conv"
+        )
+
+    # Post
+    port_batch_normalization("bn", "features.norm5")
+
+
+def convert_head(task, loader, timm_config):
+    loader.port_weight(
+        task.output_dense.kernel,
+        hf_weight_key="classifier.weight",
+        hook_fn=lambda x, _: np.transpose(np.squeeze(x)),
+    )
+    loader.port_weight(
+        task.output_dense.bias,
+        hf_weight_key="classifier.bias",
+    )

keras_hub/src/utils/timm/convert_resnet.py

@@ -1,16 +1,3 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 import numpy as np
 
 from keras_hub.src.models.resnet.resnet_backbone import ResNetBackbone

keras_hub/src/utils/timm/preset_loader.py

@@ -1,21 +1,9 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 """Convert timm models to KerasHub."""
 
 from keras_hub.src.models.image_classifier import ImageClassifier
 from keras_hub.src.utils.preset_utils import PresetLoader
 from keras_hub.src.utils.preset_utils import jax_memory_cleanup
+from keras_hub.src.utils.timm import convert_densenet
 from keras_hub.src.utils.timm import convert_resnet
 from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
 
@@ -26,6 +14,8 @@ class TimmPresetLoader(PresetLoader):
         architecture = self.config["architecture"]
         if "resnet" in architecture:
             self.converter = convert_resnet
+        if "densenet" in architecture:
+            self.converter = convert_densenet
         else:
             raise ValueError(
                 "KerasHub has no converter for timm models "

keras_hub/src/utils/transformers/__init__.py

@@ -1,13 +0,0 @@
-# Copyright 2024 The KerasHub Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.