keras-hub 0.22.1__py3-none-any.whl → 0.23.0.dev0__py3-none-any.whl

keras_hub/src/utils/timm/convert_mobilenetv5.py

@@ -0,0 +1,321 @@
+import types
+
+import keras
+import numpy as np
+
+from keras_hub.src.models.mobilenetv5.mobilenetv5_attention import (
+    MobileAttention,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_backbone import (
+    MobileNetV5Backbone,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_blocks import EdgeResidual
+from keras_hub.src.models.mobilenetv5.mobilenetv5_blocks import (
+    UniversalInvertedResidual,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_builder import (
+    convert_arch_def_to_stackwise,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_layers import ConvNormAct
+from keras_hub.src.models.mobilenetv5.mobilenetv5_layers import RmsNorm2d
+
+backbone_cls = MobileNetV5Backbone
+
+MODEL_CONFIGS = {
+    "mobilenetv5_300m": {
+        "backbone": convert_arch_def_to_stackwise(
+            [
+                # Stage 0: 128x128 in
+                [
+                    "er_r1_k3_s2_e4_c128",
+                    "er_r1_k3_s1_e4_c128",
+                    "er_r1_k3_s1_e4_c128",
+                ],
+                # Stage 1: 256x256 in
+                [
+                    "uir_r1_a3_k5_s2_e6_c256",
+                    "uir_r1_a5_k0_s1_e4_c256",
+                    "uir_r1_a3_k0_s1_e4_c256",
+                    "uir_r1_a5_k0_s1_e4_c256",
+                    "uir_r1_a3_k0_s1_e4_c256",
+                ],
+                # Stage 2: 640x640 in
+                [
+                    "uir_r1_a5_k5_s2_e6_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a5_k0_s1_e4_c640",
+                    "uir_r1_a0_k0_s1_e1_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                    "mqa_r1_k3_h12_v2_s1_d64_c640",
+                    "uir_r1_a0_k0_s1_e2_c640",
+                ],
+                # Stage 3: 1280x1280 in
+                [
+                    "uir_r1_a5_k5_s2_e6_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                    "mqa_r1_k3_h16_s1_d96_c1280",
+                    "uir_r1_a0_k0_s1_e2_c1280",
+                ],
+            ]
+        ),
+        "stem_size": 64,
+        "num_features": 2048,
+        "norm_layer": "rms_norm",
+        "act_layer": "gelu",
+        "use_msfa": True,
+        "layer_scale_init_value": 1e-5,
+    },
+}
+
+
+def convert_head(task, loader, timm_config):
+    pass
+
+
+def convert_backbone_config(timm_config):
+    timm_architecture = timm_config["architecture"]
+    if timm_architecture not in MODEL_CONFIGS:
+        raise ValueError(f"Unsupported architecture: {timm_architecture}")
+    config = MODEL_CONFIGS[timm_architecture].copy()
+    backbone_config = config.pop("backbone")
+    backbone_config.update(config)
+    return backbone_config
+
+
+def convert_weights(backbone, loader, timm_config):
+    def key_exists(key):
+        try:
+            loader.get_tensor(key)
+            return True
+        except Exception:
+            return False
+
+    def _port_weights(layer, timm_key, transpose_dims=None):
+        hf_weight_key = f"{timm_key}.weight"
+        if not key_exists(hf_weight_key):
+            return
+        hook_fn = None
+        if transpose_dims:
+
+            def transpose_hook(x, _):
+                return np.transpose(x, transpose_dims)
+
+            hook_fn = transpose_hook
+        loader.port_weight(
+            layer.kernel, hf_weight_key=hf_weight_key, hook_fn=hook_fn
+        )
+        if layer.bias is not None:
+            hf_bias_key = f"{timm_key}.bias"
+            if key_exists(hf_bias_key):
+                loader.port_weight(
+                    layer.bias,
+                    hf_weight_key=hf_bias_key,
+                )
+
+    def _port_bn(layer, timm_prefix):
+        loader.port_weight(layer.gamma, f"{timm_prefix}.weight")
+        loader.port_weight(layer.beta, f"{timm_prefix}.bias")
+        loader.port_weight(layer.moving_mean, f"{timm_prefix}.running_mean")
+        loader.port_weight(layer.moving_variance, f"{timm_prefix}.running_var")
+
+    def _port_rms_norm(layer, timm_prefix):
+        loader.port_weight(layer.gamma, f"{timm_prefix}.weight")
+
+    def _port_cna(cna_layer: ConvNormAct, timm_conv_prefix, timm_norm_prefix):
+        if isinstance(cna_layer.conv, keras.layers.DepthwiseConv2D):
+            _port_weights(
+                cna_layer.conv,
+                timm_conv_prefix,
+                transpose_dims=(2, 3, 0, 1),
+            )
+        else:
+            _port_weights(
+                cna_layer.conv,
+                timm_conv_prefix,
+                transpose_dims=(2, 3, 1, 0),
+            )
+        if key_exists(f"{timm_norm_prefix}.running_mean"):
+            _port_bn(cna_layer.norm, timm_norm_prefix)
+        else:
+            _port_rms_norm(cna_layer.norm, timm_norm_prefix)
+
+    def _port_attn(attn_layer, attn_prefix):
+        _port_weights(
+            attn_layer.query_layers[-1],
+            f"{attn_prefix}.query.proj",
+            (2, 3, 1, 0),
+        )
+        if len(attn_layer.key_layers) > 1:
+            _port_weights(
+                attn_layer.key_layers[0],
+                f"{attn_prefix}.key.down_conv",
+                (2, 3, 0, 1),
+            )
+            key_norm_layer = attn_layer.key_layers[1]
+            if isinstance(key_norm_layer, RmsNorm2d):
+                _port_rms_norm(key_norm_layer, f"{attn_prefix}.key.norm")
+            else:
+                _port_bn(key_norm_layer, f"{attn_prefix}.key.norm")
+        _port_weights(
+            attn_layer.key_layers[-1], f"{attn_prefix}.key.proj", (2, 3, 1, 0)
+        )
+        if len(attn_layer.value_layers) > 1:
+            _port_weights(
+                attn_layer.value_layers[0],
+                f"{attn_prefix}.value.down_conv",
+                (2, 3, 0, 1),
+            )
+            value_norm_layer = attn_layer.value_layers[1]
+            if isinstance(value_norm_layer, RmsNorm2d):
+                _port_rms_norm(value_norm_layer, f"{attn_prefix}.value.norm")
+            else:
+                _port_bn(value_norm_layer, f"{attn_prefix}.value.norm")
+        _port_weights(
+            attn_layer.value_layers[-1],
+            f"{attn_prefix}.value.proj",
+            (2, 3, 1, 0),
+        )
+        _port_weights(
+            attn_layer.output_proj_layers[-2],
+            f"{attn_prefix}.output.proj",
+            (2, 3, 1, 0),
+        )
+
+    stem_layer = backbone.get_layer("conv_stem")
+    _port_cna(stem_layer, "conv_stem.conv", "conv_stem.bn")
+    block_layers = [
+        layer
+        for layer in backbone.layers
+        if isinstance(
+            layer, (EdgeResidual, UniversalInvertedResidual, MobileAttention)
+        )
+    ]
+    block_counter = 0
+    for stack_idx in range(len(backbone.stackwise_num_blocks)):
+        for block_idx_in_stage in range(
+            backbone.stackwise_num_blocks[stack_idx]
+        ):
+            block = block_layers[block_counter]
+            timm_prefix = f"blocks.{stack_idx}.{block_idx_in_stage}"
+            if isinstance(block, EdgeResidual):
+                _port_cna(
+                    block.conv_exp,
+                    f"{timm_prefix}.conv_exp",
+                    f"{timm_prefix}.bn1",
+                )
+                _port_cna(
+                    block.conv_pwl,
+                    f"{timm_prefix}.conv_pwl",
+                    f"{timm_prefix}.bn2",
+                )
+            elif isinstance(block, UniversalInvertedResidual):
+                if hasattr(block, "dw_start") and not isinstance(
+                    block.dw_start, types.FunctionType
+                ):
+                    _port_cna(
+                        block.dw_start,
+                        f"{timm_prefix}.dw_start.conv",
+                        f"{timm_prefix}.dw_start.bn",
+                    )
+                _port_cna(
+                    block.pw_exp,
+                    f"{timm_prefix}.pw_exp.conv",
+                    f"{timm_prefix}.pw_exp.bn",
+                )
+                if hasattr(block, "dw_mid") and not isinstance(
+                    block.dw_mid, types.FunctionType
+                ):
+                    _port_cna(
+                        block.dw_mid,
+                        f"{timm_prefix}.dw_mid.conv",
+                        f"{timm_prefix}.dw_mid.bn",
+                    )
+                _port_cna(
+                    block.pw_proj,
+                    f"{timm_prefix}.pw_proj.conv",
+                    f"{timm_prefix}.pw_proj.bn",
+                )
+                gamma_key = f"{timm_prefix}.layer_scale.gamma"
+                if key_exists(gamma_key):
+                    loader.port_weight(block.layer_scale.gamma, gamma_key)
+            elif isinstance(block, MobileAttention):
+                _port_rms_norm(block.norm, f"{timm_prefix}.norm")
+                gamma_key = f"{timm_prefix}.layer_scale.gamma"
+                if key_exists(gamma_key):
+                    loader.port_weight(block.layer_scale.gamma, gamma_key)
+                attn_prefix = f"{timm_prefix}.attn"
+                _port_attn(block.attn, attn_prefix)
+            block_counter += 1
+    try:
+        msfa_layer = backbone.get_layer("msfa")
+        ffn = msfa_layer.ffn
+        _port_cna(ffn.pw_exp, "msfa.ffn.pw_exp.conv", "msfa.ffn.pw_exp.bn")
+        _port_cna(ffn.pw_proj, "msfa.ffn.pw_proj.conv", "msfa.ffn.pw_proj.bn")
+        _port_rms_norm(msfa_layer.norm, "msfa.norm")
+    except ValueError:
+        pass

keras_hub/src/utils/timm/preset_loader.py

@@ -7,6 +7,7 @@ from keras_hub.src.utils.timm import convert_cspnet
 from keras_hub.src.utils.timm import convert_densenet
 from keras_hub.src.utils.timm import convert_efficientnet
 from keras_hub.src.utils.timm import convert_mobilenet
+from keras_hub.src.utils.timm import convert_mobilenetv5
 from keras_hub.src.utils.timm import convert_resnet
 from keras_hub.src.utils.timm import convert_vgg
 from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
@@ -22,6 +23,8 @@ class TimmPresetLoader(PresetLoader):
             self.converter = convert_cspnet
         elif architecture.startswith("densenet"):
             self.converter = convert_densenet
+        elif architecture.startswith("mobilenetv5"):
+            self.converter = convert_mobilenetv5
         elif architecture.startswith("mobilenet"):
             self.converter = convert_mobilenet
         elif architecture.startswith("vgg"):
@@ -41,7 +44,8 @@ class TimmPresetLoader(PresetLoader):
         keras_config = self.converter.convert_backbone_config(self.config)
         backbone = cls(**{**keras_config, **kwargs})
         if load_weights:
-            jax_memory_cleanup(backbone)
+            if not self.config["architecture"].startswith("mobilenetv5"):
+                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)
@@ -54,9 +58,9 @@ class TimmPresetLoader(PresetLoader):
             )
         # Support loading the classification head for classifier models.
         kwargs["num_classes"] = self.config["num_classes"]
-        if (
-            "num_features" in self.config
-            and "mobilenet" in self.config["architecture"]
+        if "num_features" in self.config and (
+            "mobilenet" in self.config["architecture"]
+            or "mobilenetv5" in self.config["architecture"]
         ):
             kwargs["num_features"] = self.config["num_features"]
 

keras_hub/src/utils/transformers/convert_dinov2.py

@@ -29,6 +29,7 @@ def convert_backbone_config(transformers_config):
         "image_shape": (image_size, image_size, 3),
         "position_embedding_shape": (image_size, image_size),
         "antialias_in_interpolation": antialias_in_interpolation,
+        "apply_layernorm": transformers_config.get("apply_layernorm", False),
     }
 
 

keras_hub/src/utils/transformers/convert_qwen3_moe.py

@@ -0,0 +1,216 @@
+import numpy as np
+
+from keras_hub.src.models.qwen3_moe.qwen3_moe_backbone import Qwen3MoeBackbone
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = Qwen3MoeBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "head_dim": transformers_config["head_dim"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_query_heads": transformers_config["num_attention_heads"],
+        "num_key_value_heads": transformers_config["num_key_value_heads"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "moe_intermediate_dim": transformers_config["moe_intermediate_size"],
+        "num_experts": transformers_config["num_experts"],
+        "top_k": transformers_config["num_experts_per_tok"],
+        "norm_top_k_prob": transformers_config["norm_topk_prob"],
+        "decoder_sparse_step": transformers_config["decoder_sparse_step"],
+        "layer_norm_epsilon": transformers_config["rms_norm_eps"],
+        "rope_max_wavelength": transformers_config["rope_theta"],
+        "sliding_window_size": transformers_config["sliding_window"],
+        "router_aux_loss_coefficient": transformers_config[
+            "router_aux_loss_coef"
+        ],
+        "tie_word_embeddings": transformers_config.get(
+            "tie_word_embeddings", False
+        ),
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").embeddings,
+        hf_weight_key="model.embed_tokens.weight",
+    )
+    if not backbone.tie_word_embeddings:
+        loader.port_weight(
+            keras_variable=backbone.get_layer(
+                "token_embedding"
+            ).reverse_embeddings,
+            hf_weight_key="lm_head.weight",
+            # rearrange_pattern="b a -> a b",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+
+    def transpose_and_reshape(x, shape):
+        return np.reshape(np.transpose(x), shape)
+
+    for i in range(backbone.num_layers):
+        decoder_layer = backbone.get_layer(f"transformer_layer_{i}")
+
+        # Input layernorm
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layernorm.scale,
+            hf_weight_key=f"model.layers.{i}.input_layernorm.weight",
+        )
+
+        # Attention layers
+
+        ## Query
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._query_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.q_proj.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._query_dense_layer_norm.scale,
+            hf_weight_key=f"model.layers.{i}.self_attn.q_norm.weight",
+        )
+        ## Key
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._key_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.k_proj.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._key_dense_layer_norm.scale,
+            hf_weight_key=f"model.layers.{i}.self_attn.k_norm.weight",
+        )
+        ## Value
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._value_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.v_proj.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        ## Output
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._output_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.o_proj.weight",
+            # rearrange_patterns="c (a b) -> a b c",
+            # rearrange_dims={"a": backbone.num_query_heads},
+            hook_fn=transpose_and_reshape,
+        )
+
+        # MLP layers
+        if (
+            (i not in backbone.mlp_only_layers)
+            and backbone.num_experts > 0
+            and ((i + 1) % backbone.decoder_sparse_step == 0)
+        ):
+            # MoE layers
+            loader.port_weight(
+                keras_variable=decoder_layer.mlp._sparse_feedforward_gate_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.gate.weight",
+                # rearrange_patterns="b a -> a b",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            # Batched experts: gate_up_proj and down_proj
+            gate_up_proj_list = []
+            down_proj_list = []
+            for expert_idx in range(backbone.num_experts):
+                # Load gate_proj and up_proj for each expert
+                gate_proj = loader.get_tensor(
+                    f"model.layers.{i}.mlp.experts.{expert_idx}.gate_proj.weight"
+                )
+                up_proj = loader.get_tensor(
+                    f"model.layers.{i}.mlp.experts.{expert_idx}.up_proj.weight"
+                )
+                # Transpose to (hidden_dim, intermediate_dim)
+                gate_proj = np.transpose(gate_proj, axes=(1, 0))
+                up_proj = np.transpose(up_proj, axes=(1, 0))
+                # Concatenate gate_proj and up_proj along the last dimension
+                gate_up_proj = np.concatenate([gate_proj, up_proj], axis=-1)
+                gate_up_proj_list.append(gate_up_proj)
+
+                # Load down_proj for each expert
+                down_proj = loader.get_tensor(
+                    f"model.layers.{i}.mlp.experts.{expert_idx}.down_proj.weight"
+                )
+                down_proj = np.transpose(
+                    down_proj, axes=(1, 0)
+                )  # (intermediate_dim, hidden_dim)
+                down_proj_list.append(down_proj)
+
+            # Stack the lists to create batched weights
+            gate_up_proj_batched = np.stack(
+                gate_up_proj_list, axis=0
+            )  # (num_experts, hidden_dim, 2 * intermediate_dim)
+            down_proj_batched = np.stack(
+                down_proj_list, axis=0
+            )  # (num_experts, intermediate_dim, hidden_dim)
+
+            # Assign batched weights to expert_bank
+            decoder_layer.mlp.expert_bank._expert_feedforward_gate_dense.assign(
+                gate_up_proj_batched
+            )
+            decoder_layer.mlp.expert_bank._expert_feedforward_output_dense.assign(
+                down_proj_batched
+            )
+        else:
+            loader.port_weight(
+                keras_variable=decoder_layer._feedforward_intermediate_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.up_proj.weight",
+                # rearrange_patterns="b a -> a b",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            loader.port_weight(
+                keras_variable=decoder_layer._feedforward_output_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.down_proj.weight",
+                # rearrange_patterns="b a -> a b",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            loader.port_weight(
+                keras_variable=decoder_layer._feedforward_gate_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.gate_proj.weight",
+                # rearrange_patterns="b a -> a b",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+
+        # Feedforward layernorm
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_layernorm.scale,
+            hf_weight_key=f"model.layers.{i}.post_attention_layernorm.weight",
+        )
+
+    # Final normalization layer
+    loader.port_weight(
+        keras_variable=backbone.get_layer("sequence_output_layernorm").scale,
+        hf_weight_key="model.norm.weight",
+    )
+
+    return backbone
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    tokenizer_config = load_json(preset, "tokenizer.json")
+    vocab = tokenizer_config["model"]["vocab"]
+    merges = tokenizer_config["model"]["merges"]
+    merges = [" ".join(item) for item in merges]
+
+    # Load all special tokens with the exception of "reserved" ones.
+    special_tokens = set()
+    for token in tokenizer_config["added_tokens"]:
+        if not token["content"].startswith("<|reserved_special_token_"):
+            vocab[token["content"]] = token["id"]
+            special_tokens.add(token["content"])
+
+    kwargs.update(
+        {
+            "unsplittable_tokens": list(special_tokens),
+        }
+    )
+
+    return cls(vocabulary=vocab, merges=merges, **kwargs)