keras-hub-nightly 0.16.1.dev202409260340__py3-none-any.whl → 0.16.1.dev202409280337__py3-none-any.whl

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
@@ -322,3 +308,109 @@ def any_equal(inputs, values, padding_mask):
         output = ops.logical_or(output, value_equality)
 
     return ops.logical_and(output, padding_mask)
+
+
+def target_gather(
+    targets,
+    indices,
+    mask=None,
+    mask_val=0.0,
+):
+    """A utility function wrapping `ops.take`, which deals with:
+        1) both batched and unbatched `targets`.
+        2) when unbatched `targets` have empty rows, the result will be filled
+            with `mask_val`.
+        3) target masking.
+
+    Args:
+        targets: `[N, ...]` or `[batch_size, N, ...]` Tensor representing
+            targets such as boxes, keypoints, etc.
+        indices: [M] or [batch_size, M] int32 Tensor representing indices within
+            `targets` to gather.
+        mask: `[M, ...]` or `[batch_size, M, ...]` boolean Tensor
+            representing the masking for each target. `True` means the
+            corresponding entity should be masked to `mask_val`, `False`
+            means the corresponding entity should be the target value.
+            Defaults to `None`.
+        mask_val: float. representing the masking value if `mask` is True
+            on the entity.
+            Defaults to `0.0`
+
+    Returns:
+        targets: `[M, ...]` or `[batch_size, M, ...]` Tensor representing
+            selected targets.
+
+        Raise:
+            ValueError: If `targets` is higher than rank 3.
+    """
+    targets_shape = list(targets.shape)
+    if len(targets_shape) > 3:
+        raise ValueError(
+            f"`target_gather` does not support `targets` with rank "
+            f"larger than 3, got {len(targets.shape)}"
+        )
+
+    def gather_unbatched(labels, match_indices, mask, mask_val):
+        """Gather based on unbatched labels and boxes."""
+        num_gt_boxes = labels.shape[0]
+
+        def assign_when_rows_empty():
+            if len(labels.shape) > 1:
+                mask_shape = [match_indices.shape[0], labels.shape[-1]]
+            else:
+                mask_shape = [match_indices.shape[0]]
+            return ops.cast(mask_val, labels.dtype) * ops.ones(
+                mask_shape, dtype=labels.dtype
+            )
+
+        def assign_when_rows_not_empty():
+            targets = ops.take(labels, match_indices, axis=0)
+            if mask is None:
+                return targets
+            else:
+                masked_targets = ops.cast(
+                    mask_val, labels.dtype
+                ) * ops.ones_like(mask, dtype=labels.dtype)
+                return ops.where(mask, masked_targets, targets)
+
+        if num_gt_boxes > 0:
+            return assign_when_rows_not_empty()
+        else:
+            return assign_when_rows_empty()
+
+    def _gather_batched(labels, match_indices, mask, mask_val):
+        """Gather based on batched labels."""
+        batch_size = labels.shape[0]
+        if batch_size == 1:
+            if mask is not None:
+                result = gather_unbatched(
+                    ops.squeeze(labels, axis=0),
+                    ops.squeeze(match_indices, axis=0),
+                    ops.squeeze(mask, axis=0),
+                    mask_val,
+                )
+            else:
+                result = gather_unbatched(
+                    ops.squeeze(labels, axis=0),
+                    ops.squeeze(match_indices, axis=0),
+                    None,
+                    mask_val,
+                )
+            return ops.expand_dims(result, axis=0)
+        else:
+            targets = ops.take_along_axis(
+                labels, ops.expand_dims(match_indices, axis=-1), axis=1
+            )
+
+            if mask is None:
+                return targets
+            else:
+                masked_targets = ops.cast(
+                    mask_val, labels.dtype
+                ) * ops.ones_like(mask, dtype=labels.dtype)
+                return ops.where(mask, masked_targets, targets)
+
+    if len(targets_shape) <= 2:
+        return gather_unbatched(targets, indices, mask, mask_val)
+    elif len(targets_shape) == 3:
+        return _gather_batched(targets, indices, mask, mask_val)

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

@@ -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.densenet.densenet_backbone import DenseNetBackbone

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,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.
 """Convert timm models to KerasHub."""
 
 from keras_hub.src.models.image_classifier import ImageClassifier

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.

keras_hub/src/utils/transformers/convert_albert.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.albert.albert_backbone import AlbertBackbone

keras_hub/src/utils/transformers/convert_bart.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.bart.bart_backbone import BartBackbone

keras_hub/src/utils/transformers/convert_bert.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.bert.bert_backbone import BertBackbone

keras_hub/src/utils/transformers/convert_distilbert.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.distil_bert.distil_bert_backbone import (

keras_hub/src/utils/transformers/convert_gemma.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.gemma.gemma_backbone import GemmaBackbone

keras_hub/src/utils/transformers/convert_gpt2.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.gpt2.gpt2_backbone import GPT2Backbone

keras_hub/src/utils/transformers/convert_llama3.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.llama3.llama3_backbone import Llama3Backbone

keras_hub/src/utils/transformers/convert_mistral.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.mistral.mistral_backbone import MistralBackbone

keras_hub/src/utils/transformers/convert_pali_gemma.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.pali_gemma.pali_gemma_backbone import (

keras_hub/src/utils/transformers/preset_loader.py

@@ -1,19 +1,5 @@
-# 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 huggingface models to KerasHub."""
 
-
 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.transformers import convert_albert
@@ -69,7 +55,7 @@ class TransformersPresetLoader(PresetLoader):
                 self.converter.convert_weights(backbone, loader, self.config)
         return backbone
 
-    def load_tokenizer(self, cls, **kwargs):
+    def load_tokenizer(self, cls, config_name="tokenizer.json", **kwargs):
         return self.converter.convert_tokenizer(cls, self.preset, **kwargs)
 
     def load_image_converter(self, cls, **kwargs):

keras_hub/src/utils/transformers/safetensor_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.
 import contextlib
 
 from keras_hub.src.utils.preset_utils import SAFETENSOR_CONFIG_FILE
@@ -26,7 +13,7 @@ except ImportError:
 
 
 class SafetensorLoader(contextlib.ExitStack):
-    def __init__(self, preset, prefix=None):
+    def __init__(self, preset, prefix=None, fname=None):
         super().__init__()
 
         if safetensors is None:
@@ -44,6 +31,13 @@ class SafetensorLoader(contextlib.ExitStack):
         self.safetensor_files = {}
         self.prefix = prefix
 
+        if fname is not None and self.safetensor_config is not None:
+            raise ValueError(
+                f"Cannot specify `fname` if {SAFETENSOR_CONFIG_FILE} exists. "
+                f"Received: fname={fname}"
+            )
+        self.fname = fname  # Specify the name of the safetensor file.
+
     def get_prefixed_key(self, hf_weight_key, dict_like):
         """
         Determine and return a prefixed key for a given hf weight key.
@@ -71,7 +65,7 @@ class SafetensorLoader(contextlib.ExitStack):
 
     def get_tensor(self, hf_weight_key):
         if self.safetensor_config is None:
-            fname = SAFETENSOR_FILE
+            fname = self.fname if self.fname is not None else SAFETENSOR_FILE
         else:
             full_key = self.get_prefixed_key(
                 hf_weight_key, self.safetensor_config["weight_map"]