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

keras_hub/src/utils/tensor_utils.py

@@ -12,7 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
+import functools
+import inspect
+import threading
+
 import keras
+import numpy as np
 from keras import ops
 
 try:
@@ -23,6 +29,181 @@ except ImportError:
     tf_text = None
 
 
+NO_CONVERT_COUNTER = threading.local()
+
+
+@contextlib.contextmanager
+def no_convert_scope():
+    try:
+        NO_CONVERT_COUNTER.count = getattr(NO_CONVERT_COUNTER, "count", 0) + 1
+        yield
+    finally:
+        NO_CONVERT_COUNTER.count = getattr(NO_CONVERT_COUNTER, "count", 0) - 1
+
+
+def in_no_convert_scope():
+    return getattr(NO_CONVERT_COUNTER, "count", 0) > 0
+
+
+def preprocessing_function(fn):
+    """Wraps a preprocessing function to handle tf tensor conversion."""
+    if tf is None:
+        return fn
+
+    params = inspect.signature(fn).parameters
+    accepts_labels = all(k in params for k in ("x", "y", "sample_weight"))
+    if not accepts_labels:
+
+        @functools.wraps(fn)
+        def wrapper(self, x, **kwargs):
+            with tf.device("cpu"):
+                x = convert_preprocessing_inputs(x)
+                with no_convert_scope():
+                    x = fn(self, x, **kwargs)
+                return convert_preprocessing_outputs(x)
+
+    else:
+
+        @functools.wraps(fn)
+        def wrapper(self, x, y=None, sample_weight=None, **kwargs):
+            with tf.device("cpu"):
+                x, y, sample_weight = convert_preprocessing_inputs(
+                    (x, y, sample_weight)
+                )
+                with no_convert_scope():
+                    x = fn(self, x, y=y, sample_weight=sample_weight, **kwargs)
+                return convert_preprocessing_outputs(x)
+
+    return wrapper
+
+
+def convert_preprocessing_inputs(x):
+    """Convert raw inputs for preprocessing.
+
+    This function is used to convert raw inputs (strings, lists, `np.ndarray`s,
+    `jax.Array`s, `torch.Tensor`s, etc) to a canonical format for
+    preprocessing layers. All inputs will be converted to backend tensors if
+    possible, except ragged inputs and string inputs which be converted to tf
+    tensors regardless of backend.
+
+    `tuple` and `list` elements are handled differently by this function. A
+    `tuple` is assumed to enumerate separate inputs, and a `list` is assumed to
+    enumerate elements in a single array-like input. This makes it possible to
+    represent ragged and string inputs in a multi-backend format, as shown in
+    the examples below.
+
+    Examples:
+    ```python
+    # Two ragged arrays of token ids.
+    x = ([[1, 2, 3], [4, 5]], [[1, 2], [3, 4, 5]])
+    keras_hub.utils.convert_preprocessing_inputs(x)
+
+    # A batch of three samples each with two string segments.
+    x = (["hi", "hello", "hey"], ["bye", "later", "so long"])
+    keras_hub.utils.convert_preprocessing_inputs(x)
+
+    # A batch of features in a dictionary.
+    x = {
+        "text": ["hi", "hello", "hey"],
+        "images": np.ones((3, 64, 64, 3)),
+        "labels": [1, 0, 1],
+    }
+    keras_hub.utils.convert_preprocessing_inputs(x)
+    ```
+    """
+    if not tf.executing_eagerly() or in_no_convert_scope():
+        return x
+
+    if isinstance(x, dict):
+        return {k: convert_preprocessing_inputs(x[k]) for k, v in x.items()}
+    if isinstance(x, tuple):
+        return tuple(convert_preprocessing_inputs(v) for v in x)
+    if isinstance(x, (str, bytes)):
+        return tf.constant(x)
+    if isinstance(x, list):
+        try:
+            numpy_x = np.array(x)
+        except ValueError as e:
+            # If numpy conversion failed, try converting to a ragged array.
+            try:
+                return tf.ragged.constant(x)
+            except ValueError:
+                # If ragged conversion failed return to the numpy error.
+                raise e
+        # If we have a string input, use tf.tensor.
+        if numpy_x.dtype.type is np.str_ or numpy_x.dtype.type is np.bytes_:
+            return tf.convert_to_tensor(x)
+        # Numpy will default to int64, int32 works with more ops.
+        if numpy_x.dtype == np.int64:
+            numpy_x = numpy_x.astype(np.int32)
+        # We have non-ragged, non-string input. Use backbend type.
+        x = ops.convert_to_tensor(numpy_x)
+        # Torch will complain about device placement for GPU tensors.
+        if keras.config.backend() == "torch":
+            x = x.cpu()
+        return x
+    if is_tensor_type(x):
+        # String or ragged types we keep as tf.
+        if isinstance(x, tf.RaggedTensor) or x.dtype == tf.string:
+            return x
+        # If we have a string input, use tf.tensor.
+        if isinstance(x, np.ndarray) and x.dtype.type is np.str_:
+            return tf.convert_to_tensor(x)
+        x = ops.convert_to_tensor(x)
+        # Torch will complain about device placement for GPU tensors.
+        if keras.config.backend() == "torch":
+            x = x.cpu()
+        return x
+    return x
+
+
+def convert_preprocessing_outputs(x):
+    """Convert outputs after preprocessing to a backend agnostic format.
+
+    This function is used to convert `tf.Tensor` and `tf.RaggedTensor` output
+    from preprocessing layers to either:
+
+    - The correct tensor type for the Keras backend framework.
+    - Python lists, in the case of ragged and string data.
+
+    This will automatically be called when on the output of preprocessing
+    layers or `keras_hub.models.Task`s with preprocessing included. It could be
+    used directly to convert a `tf.data.Dataset` output to a backend agnostic
+    type.
+
+    Examples:
+    ```python
+    # Two ragged arrays of token ids.
+    x = tf.ragged.constant([[1, 2, 3], [4, 5]])
+    keras_hub.utils.convert_preprocessing_outputs(x)
+
+    # A batch of three samples each with two string segments.
+    x = (tf.constant["hi", "yo", "hey"]), tf.constant(["bye", "ciao", ""]))
+    keras_hub.utils.convert_preprocessing_outputs(x)
+
+    # A batch of features in a dictionary.
+    x = {
+        "text": tf.constant(["hi", "hello", "hey"]),
+        "images": tf.ones((3, 64, 64, 3)),
+        "labels": tf.constant([1, 0, 1]),
+    }
+    keras_hub.utils.convert_preprocessing_outputs(x)
+    ```
+    """
+    if not tf.executing_eagerly() or in_no_convert_scope():
+        return x
+
+    def convert(x):
+        if x is None:
+            return x
+        if isinstance(x, tf.RaggedTensor) or x.dtype == tf.string:
+            return tensor_to_list(x)
+        dtype = keras.backend.standardize_dtype(x.dtype)
+        return ops.convert_to_tensor(x, dtype=dtype)
+
+    return keras.tree.map_structure(convert, x)
+
+
 def _decode_strings_to_utf8(inputs):
     """Recursively decodes to list of strings with 'utf-8' encoding."""
     if isinstance(inputs, bytes):
@@ -52,75 +233,15 @@ def tensor_to_list(inputs):
     return list_outputs
 
 
-def convert_to_backend_tensor_or_python_list(x):
-    """
-    Convert a tensor to the backend friendly representation of the data.
-
-    This wraps `ops.convert_to_tensor` to account for the fact that torch and
-    jax both lack native types for ragged and string data.
-
-    If we encounter one of these types in torch or jax, we will instead covert
-    the tensor to simple pythonic types (lists of strings).
-    """
-    if isinstance(x, tf.RaggedTensor) or getattr(x, "dtype", None) == tf.string:
-        return tensor_to_list(x)
-    dtype = getattr(x, "dtype", "float32")
-    dtype = keras.backend.standardize_dtype(dtype)
-    return ops.convert_to_tensor(x, dtype=dtype)
-
-
 def convert_to_ragged_batch(inputs):
-    """Convert pythonic or numpy-like input to a 2-D `tf.RaggedTensor`.
-
-    This is useful for text preprocessing layers which deal with already
-    tokenized or split text.
-
-    Args:
-        inputs: A pythonic or numpy-like input to covert. This input should
-            represent a possibly batched list of token sequences.
-
-    Returns:
-        An `(inputs, unbatched, rectangular)` tuple, where `inputs` is a
-        2-D `tf.RaggedTensor`, `unbatched` is `True` if the inputs were
-        origianlly rank 1, and `rectangular` is `True` if the inputs rows are
-        all of equal lengths.
-    """
-    # `tf.keras.layers.Layer` does a weird conversion in __call__, where a list
-    # of lists of ints will become a list of list of scalar tensors. We could
-    # clean this up if we no longer need to care about that case.
-    if isinstance(inputs, (list, tuple)):
-        if isinstance(inputs[0], (list, tuple)):
-            rectangular = len(set([len(row) for row in inputs])) == 1
-            rows = [
-                tf.convert_to_tensor(row, dtype_hint="int32") for row in inputs
-            ]
-            inputs = tf.ragged.stack(rows).with_row_splits_dtype("int64")
-        else:
-            inputs = tf.convert_to_tensor(inputs)
-            rectangular = True
-    elif isinstance(inputs, tf.Tensor):
-        rectangular = True
-    elif isinstance(inputs, tf.RaggedTensor):
-        rectangular = False
-    elif hasattr(inputs, "__array__"):
-        inputs = tf.convert_to_tensor(ops.convert_to_numpy(inputs))
-        rectangular = True
-    else:
-        raise ValueError(
-            f"Unknown tensor type. Tensor input can be passed as "
-            "tensors, numpy arrays, or python lists. Received: "
-            f"`type(inputs)={type(inputs)}`"
-        )
-    if inputs.shape.rank < 1 or inputs.shape.rank > 2:
-        raise ValueError(
-            f"Tokenized tensor input should be rank 1 (unbatched) or "
-            f"rank 2 (batched). Received: `inputs.shape={input.shape}`"
-        )
+    """Ensure a tf.Tensor is a ragged rank 2 tensor."""
+    if not isinstance(inputs, (tf.RaggedTensor, tf.Tensor)):
+        inputs = tf.convert_to_tensor(inputs)
     unbatched = inputs.shape.rank == 1
-    rectangular = rectangular or unbatched
+    rectangular = isinstance(inputs, tf.Tensor)
     if unbatched:
         inputs = tf.expand_dims(inputs, 0)
-    if isinstance(inputs, tf.Tensor):
+    if rectangular:
         inputs = tf.RaggedTensor.from_tensor(inputs)
     return inputs, unbatched, rectangular
 
@@ -135,10 +256,7 @@ def truncate_at_token(inputs, token, mask):
 
 def strip_to_ragged(token_ids, mask, ids_to_strip):
     """Remove masked and special tokens from a sequence before detokenizing."""
-    token_ids = ops.convert_to_numpy(token_ids)
-    token_ids = token_ids.astype("int32")
-    mask = ops.convert_to_numpy(mask)
-    mask = mask.astype("bool")
+    mask = tf.cast(mask, "bool")
     for id in ids_to_strip:
         mask = mask & (token_ids != id)
     return tf.ragged.boolean_mask(token_ids, mask)

keras_hub/src/utils/timm/convert_resnet.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+from keras_hub.src.models.resnet.resnet_backbone import ResNetBackbone
+
+backbone_cls = ResNetBackbone
 
 
 def convert_backbone_config(timm_config):
@@ -56,6 +55,8 @@ def convert_backbone_config(timm_config):
         stackwise_num_strides=[1, 2, 2, 2],
         block_type=block_type,
         use_pre_activation=use_pre_activation,
+        input_conv_filters=[64],
+        input_conv_kernel_sizes=[7],
     )
 
 
@@ -100,10 +101,10 @@ def convert_weights(backbone, loader, timm_config):
     for stack_index in range(num_stacks):
         for block_idx in range(backbone.stackwise_num_blocks[stack_index]):
             if version == "v1":
-                keras_name = f"v1_stack{stack_index}_block{block_idx}"
+                keras_name = f"stack{stack_index}_block{block_idx}"
                 hf_name = f"layer{stack_index+1}.{block_idx}"
             else:
-                keras_name = f"v2_stack{stack_index}_block{block_idx}"
+                keras_name = f"stack{stack_index}_block{block_idx}"
                 hf_name = f"stages.{stack_index}.blocks.{block_idx}"
 
             if version == "v1":
@@ -159,13 +160,15 @@ def convert_weights(backbone, loader, timm_config):
     normalization_layer.build(normalization_layer._build_input_shape)
 
 
-def load_resnet_backbone(cls, preset, load_weights, **kwargs):
-    timm_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(timm_config)
-    backbone = cls(**keras_config, **kwargs)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        # Use prefix="" to avoid using `get_prefixed_key`.
-        with SafetensorLoader(preset, prefix="") as loader:
-            convert_weights(backbone, loader, timm_config)
-    return backbone
+def convert_head(task, loader, timm_config):
+    v2 = "resnetv2_" in timm_config["architecture"]
+    prefix = "head.fc." if v2 else "fc."
+    loader.port_weight(
+        task.output_dense.kernel,
+        hf_weight_key=prefix + "weight",
+        hook_fn=lambda x, _: np.transpose(np.squeeze(x)),
+    )
+    loader.port_weight(
+        task.output_dense.bias,
+        hf_weight_key=prefix + "bias",
+    )

keras_hub/src/utils/timm/preset_loader.py

@@ -0,0 +1,66 @@
+# 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_resnet
+from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+
+
+class TimmPresetLoader(PresetLoader):
+    def __init__(self, preset, config):
+        super().__init__(preset, config)
+        architecture = self.config["architecture"]
+        if "resnet" in architecture:
+            self.converter = convert_resnet
+        else:
+            raise ValueError(
+                "KerasHub has no converter for timm models "
+                f"with architecture `'{architecture}'`."
+            )
+
+    def check_backbone_class(self):
+        return self.converter.backbone_cls
+
+    def load_backbone(self, cls, load_weights, **kwargs):
+        keras_config = self.converter.convert_backbone_config(self.config)
+        backbone = cls(**{**keras_config, **kwargs})
+        if load_weights:
+            jax_memory_cleanup(backbone)
+            # Use prefix="" to avoid using `get_prefixed_key`.
+            with SafetensorLoader(self.preset, prefix="") as loader:
+                self.converter.convert_weights(backbone, loader, self.config)
+        return backbone
+
+    def load_task(self, cls, load_weights, load_task_weights, **kwargs):
+        if not load_task_weights or not issubclass(cls, ImageClassifier):
+            return super().load_task(
+                cls, load_weights, load_task_weights, **kwargs
+            )
+        # Support loading the classification head for classifier models.
+        kwargs["num_classes"] = self.config["num_classes"]
+        task = super().load_task(cls, load_weights, load_task_weights, **kwargs)
+        if load_task_weights:
+            with SafetensorLoader(self.preset, prefix="") as loader:
+                self.converter.convert_head(task, loader, self.config)
+        return task
+
+    def load_image_converter(self, cls, **kwargs):
+        pretrained_cfg = self.config.get("pretrained_cfg", None)
+        if not pretrained_cfg or "input_size" not in pretrained_cfg:
+            return None
+        input_size = pretrained_cfg["input_size"]
+        return cls(width=input_size[1], height=input_size[2])

keras_hub/src/utils/transformers/convert_albert.py

@@ -0,0 +1,193 @@
+# 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
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = AlbertBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "embedding_dim": transformers_config["embedding_size"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "num_groups": transformers_config["num_hidden_groups"],
+        "num_inner_repetitions": transformers_config["inner_group_num"],
+        "dropout": transformers_config["attention_probs_dropout_prob"],
+        "max_sequence_length": transformers_config["max_position_embeddings"],
+        "num_segments": transformers_config["type_vocab_size"],
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embeddings
+    loader.port_weight(
+        keras_variable=backbone.token_embedding.embeddings,
+        hf_weight_key="albert.embeddings.word_embeddings.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.position_embedding.position_embeddings,
+        hf_weight_key="albert.embeddings.position_embeddings.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.segment_embedding.embeddings,
+        hf_weight_key="albert.embeddings.token_type_embeddings.weight",
+    )
+
+    # Normalization
+    loader.port_weight(
+        keras_variable=backbone.embeddings_layer_norm.gamma,
+        hf_weight_key="albert.embeddings.LayerNorm.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings_layer_norm.beta,
+        hf_weight_key="albert.embeddings.LayerNorm.bias",
+    )
+
+    # Encoder Embeddings
+    loader.port_weight(
+        keras_variable=backbone.embeddings_projection.kernel,
+        hf_weight_key="albert.encoder.embedding_hidden_mapping_in.weight",
+        hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings_projection.bias,
+        hf_weight_key="albert.encoder.embedding_hidden_mapping_in.bias",
+    )
+
+    # Encoder Group Layers
+    for group_idx in range(backbone.num_groups):
+        for inner_layer_idx in range(backbone.num_inner_repetitions):
+            keras_group = backbone.get_layer(
+                f"group_{group_idx}_inner_layer_{inner_layer_idx}"
+            )
+            hf_group_prefix = (
+                "albert.encoder.albert_layer_groups."
+                f"{group_idx}.albert_layers.{inner_layer_idx}."
+            )
+
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.query_dense.kernel,
+                hf_weight_key=f"{hf_group_prefix}attention.query.weight",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    np.transpose(hf_tensor), keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.query_dense.bias,
+                hf_weight_key=f"{hf_group_prefix}attention.query.bias",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    hf_tensor, keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.key_dense.kernel,
+                hf_weight_key=f"{hf_group_prefix}attention.key.weight",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    np.transpose(hf_tensor), keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.key_dense.bias,
+                hf_weight_key=f"{hf_group_prefix}attention.key.bias",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    hf_tensor, keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.value_dense.kernel,
+                hf_weight_key=f"{hf_group_prefix}attention.value.weight",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    np.transpose(hf_tensor), keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.value_dense.bias,
+                hf_weight_key=f"{hf_group_prefix}attention.value.bias",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    hf_tensor, keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.output_dense.kernel,
+                hf_weight_key=f"{hf_group_prefix}attention.dense.weight",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    np.transpose(hf_tensor), keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer.output_dense.bias,
+                hf_weight_key=f"{hf_group_prefix}attention.dense.bias",
+                hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                    hf_tensor, keras_shape
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer_norm.gamma,
+                hf_weight_key=f"{hf_group_prefix}attention.LayerNorm.weight",
+            )
+            loader.port_weight(
+                keras_variable=keras_group._self_attention_layer_norm.beta,
+                hf_weight_key=f"{hf_group_prefix}attention.LayerNorm.bias",
+            )
+            loader.port_weight(
+                keras_variable=keras_group._feedforward_intermediate_dense.kernel,
+                hf_weight_key=f"{hf_group_prefix}ffn.weight",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._feedforward_intermediate_dense.bias,
+                hf_weight_key=f"{hf_group_prefix}ffn.bias",
+            )
+            loader.port_weight(
+                keras_variable=keras_group._feedforward_output_dense.kernel,
+                hf_weight_key=f"{hf_group_prefix}ffn_output.weight",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            loader.port_weight(
+                keras_variable=keras_group._feedforward_output_dense.bias,
+                hf_weight_key=f"{hf_group_prefix}ffn_output.bias",
+            )
+            loader.port_weight(
+                keras_variable=keras_group._feedforward_layer_norm.gamma,
+                hf_weight_key=f"{hf_group_prefix}full_layer_layer_norm.weight",
+            )
+            loader.port_weight(
+                keras_variable=keras_group._feedforward_layer_norm.beta,
+                hf_weight_key=f"{hf_group_prefix}full_layer_layer_norm.bias",
+            )
+
+    # Pooler
+    loader.port_weight(
+        keras_variable=backbone.pooled_dense.kernel,
+        hf_weight_key="albert.pooler.weight",
+        hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+    )
+    loader.port_weight(
+        keras_variable=backbone.pooled_dense.bias,
+        hf_weight_key="albert.pooler.bias",
+    )
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    return cls(get_file(preset, "spiece.model"), **kwargs)