keras-hub 0.21.1__py3-none-any.whl → 0.22.0.dev0__py3-none-any.whl

keras_hub/src/utils/transformers/convert_deit.py

@@ -0,0 +1,155 @@
+import numpy as np
+
+from keras_hub.src.models.deit.deit_backbone import DeiTBackbone
+
+backbone_cls = DeiTBackbone
+
+
+def convert_backbone_config(transformers_config):
+    image_size = transformers_config["image_size"]
+    return {
+        "image_shape": (image_size, image_size, 3),
+        "patch_size": transformers_config["patch_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "dropout_rate": transformers_config["hidden_dropout_prob"],
+        "attention_dropout": transformers_config[
+            "attention_probs_dropout_prob"
+        ],
+        "layer_norm_epsilon": transformers_config["layer_norm_eps"],
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    def port_ln(keras_variable, weight_key):
+        loader.port_weight(keras_variable.gamma, f"{weight_key}.weight")
+        loader.port_weight(keras_variable.beta, f"{weight_key}.bias")
+
+    def port_dense(keras_variable, weight_key):
+        loader.port_weight(
+            keras_variable.kernel,
+            f"{weight_key}.weight",
+            hook_fn=lambda x, _: x.T,
+        )
+        if keras_variable.bias is not None:
+            loader.port_weight(keras_variable.bias, f"{weight_key}.bias")
+
+    def port_mha(keras_variable, weight_key, num_heads, hidden_dim):
+        # query
+        loader.port_weight(
+            keras_variable.query_dense.kernel,
+            f"{weight_key}.attention.query.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.query_dense.bias,
+            f"{weight_key}.attention.query.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # key
+        loader.port_weight(
+            keras_variable.key_dense.kernel,
+            f"{weight_key}.attention.key.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.key_dense.bias,
+            f"{weight_key}.attention.key.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # value
+        loader.port_weight(
+            keras_variable.value_dense.kernel,
+            f"{weight_key}.attention.value.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.value_dense.bias,
+            f"{weight_key}.attention.value.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # output
+        loader.port_weight(
+            keras_variable.output_dense.kernel,
+            f"{weight_key}.output.dense.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (num_heads, hidden_dim // num_heads, hidden_dim)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.output_dense.bias, f"{weight_key}.output.dense.bias"
+        )
+
+    loader.port_weight(
+        keras_variable=backbone.layers[1].patch_embedding.kernel,
+        hf_weight_key="deit.embeddings.patch_embeddings.projection.weight",
+        hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
+    )
+
+    loader.port_weight(
+        backbone.layers[1].patch_embedding.bias,
+        "deit.embeddings.patch_embeddings.projection.bias",
+    )
+
+    loader.port_weight(
+        backbone.layers[1].class_token,
+        "deit.embeddings.cls_token",
+    )
+
+    loader.port_weight(
+        backbone.layers[1].distillation_token,
+        "deit.embeddings.distillation_token",
+    )
+
+    loader.port_weight(
+        backbone.layers[1].position_embedding,
+        "deit.embeddings.position_embeddings",
+    )
+
+    encoder_layers = backbone.layers[2].encoder_layers
+    for i, encoder_block in enumerate(encoder_layers):
+        prefix = "deit.encoder.layer"
+        num_heads = encoder_block.num_heads
+        hidden_dim = encoder_block.hidden_dim
+
+        port_mha(
+            encoder_block.mha,
+            f"{prefix}.{i}.attention",
+            num_heads,
+            hidden_dim,
+        )
+        port_ln(encoder_block.layer_norm_1, f"{prefix}.{i}.layernorm_before")
+        port_ln(encoder_block.layer_norm_2, f"{prefix}.{i}.layernorm_after")
+
+        port_dense(encoder_block.mlp.dense, f"{prefix}.{i}.intermediate.dense")
+        port_dense(
+            encoder_block.output_layer.dense, f"{prefix}.{i}.output.dense"
+        )
+    port_ln(backbone.layers[2].layer_norm, "deit.layernorm")
+
+
+def convert_head(task, loader, transformers_config):
+    prefix = "cls_classifier."
+    loader.port_weight(
+        task.output_dense.kernel,
+        hf_weight_key=prefix + "weight",
+        hook_fn=lambda x, _: x.T,
+    )
+    loader.port_weight(
+        task.output_dense.bias,
+        hf_weight_key=prefix + "bias",
+    )

keras_hub/src/utils/transformers/convert_dinov2.py

@@ -0,0 +1,180 @@
+import numpy as np
+
+from keras_hub.src.models.dinov2.dinov2_backbone import DINOV2Backbone
+
+backbone_cls = DINOV2Backbone
+
+
+def convert_backbone_config(transformers_config):
+    model_type = transformers_config["model_type"]
+    antialias_in_interpolation = False if model_type == "dinov2" else True
+    image_size = transformers_config["image_size"]
+    intermediate_dim = int(
+        transformers_config["hidden_size"] * transformers_config["mlp_ratio"]
+    )
+    return {
+        "patch_size": transformers_config["patch_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "intermediate_dim": intermediate_dim,
+        "layer_scale_init_value": transformers_config["layerscale_value"],
+        "num_register_tokens": transformers_config.get(
+            "num_register_tokens", 0
+        ),
+        "use_mask_token": transformers_config.get("use_mask_token", True),
+        "use_swiglu_ffn": transformers_config["use_swiglu_ffn"],
+        "dropout_rate": transformers_config["hidden_dropout_prob"],
+        "drop_path_rate": transformers_config["drop_path_rate"],
+        "image_shape": (image_size, image_size, 3),
+        "position_embedding_shape": (image_size, image_size),
+        "antialias_in_interpolation": antialias_in_interpolation,
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    if not isinstance(backbone, DINOV2Backbone):
+        raise ValueError(
+            "The provided backbone must be an instance of DINOV2Backbone. "
+            f"Received: {type(backbone)}"
+        )
+
+    def port_ln(keras_variable, weight_key):
+        loader.port_weight(keras_variable.gamma, f"{weight_key}.weight")
+        loader.port_weight(keras_variable.beta, f"{weight_key}.bias")
+
+    def port_dense(keras_variable, weight_key):
+        loader.port_weight(
+            keras_variable.kernel,
+            f"{weight_key}.weight",
+            hook_fn=lambda x, _: x.T,
+        )
+        if keras_variable.bias is not None:
+            loader.port_weight(keras_variable.bias, f"{weight_key}.bias")
+
+    def port_mha(keras_variable, weight_key, num_heads, hidden_dim):
+        # query
+        loader.port_weight(
+            keras_variable.query_dense.kernel,
+            f"{weight_key}.attention.query.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.query_dense.bias,
+            f"{weight_key}.attention.query.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # key
+        loader.port_weight(
+            keras_variable.key_dense.kernel,
+            f"{weight_key}.attention.key.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.key_dense.bias,
+            f"{weight_key}.attention.key.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # value
+        loader.port_weight(
+            keras_variable.value_dense.kernel,
+            f"{weight_key}.attention.value.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.value_dense.bias,
+            f"{weight_key}.attention.value.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # output
+        loader.port_weight(
+            keras_variable.output_dense.kernel,
+            f"{weight_key}.output.dense.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (num_heads, hidden_dim // num_heads, hidden_dim)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.output_dense.bias, f"{weight_key}.output.dense.bias"
+        )
+
+    # Embedding.
+    loader.port_weight(
+        keras_variable=backbone.embeddings.cls_token,
+        hf_weight_key="embeddings.cls_token",
+    )
+    if backbone.use_mask_token:
+        loader.port_weight(
+            keras_variable=backbone.embeddings.mask_token,
+            hf_weight_key="embeddings.mask_token",
+        )
+    if backbone.num_register_tokens > 0:
+        loader.port_weight(
+            keras_variable=backbone.embeddings.register_tokens,
+            hf_weight_key="embeddings.register_tokens",
+        )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.position_embeddings,
+        hf_weight_key="embeddings.position_embeddings",
+    )
+    # Interpolate position embeddings to match the image shape.
+    backbone.embeddings.interpolated_position_embeddings.assign(
+        backbone.embeddings._interpolate_position_embeddings(
+            backbone.embeddings.position_embeddings,
+            patch_size=backbone.patch_size,
+            source_shape=backbone.embeddings.position_embedding_shape,
+            target_shape=backbone.image_shape,
+            antialias=backbone.embeddings.antialias_in_interpolation,
+        )
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.patch_embeddings.projection.kernel,
+        hf_weight_key="embeddings.patch_embeddings.projection.weight",
+        hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.patch_embeddings.projection.bias,
+        hf_weight_key="embeddings.patch_embeddings.projection.bias",
+    )
+
+    # Encoder.
+    hidden_dim = backbone.hidden_dim
+    num_heads = backbone.num_heads
+    for i, layer in enumerate(backbone.encoder.layers):
+        prefix = f"encoder.layer.{i}"
+        port_ln(layer.norm1, f"{prefix}.norm1")
+        port_mha(
+            layer.attention.attention,
+            f"{prefix}.attention",
+            num_heads,
+            hidden_dim,
+        )
+        loader.port_weight(
+            keras_variable=layer.layer_scale1.lambda1,
+            hf_weight_key=f"{prefix}.layer_scale1.lambda1",
+        )
+        port_ln(layer.norm2, f"{prefix}.norm2")
+        if backbone.use_swiglu_ffn:
+            port_dense(layer.mlp.weights_in, f"{prefix}.mlp.weights_in")
+            port_dense(layer.mlp.weights_out, f"{prefix}.mlp.weights_out")
+        else:
+            port_dense(layer.mlp.fc1, f"{prefix}.mlp.fc1")
+            port_dense(layer.mlp.fc2, f"{prefix}.mlp.fc2")
+        loader.port_weight(
+            keras_variable=layer.layer_scale2.lambda1,
+            hf_weight_key=f"{prefix}.layer_scale2.lambda1",
+        )
+
+    port_ln(backbone.layernorm, "layernorm")

keras_hub/src/utils/transformers/convert_esm.py

@@ -0,0 +1,159 @@
+import numpy as np
+
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = ESMBackbone
+
+
+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"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "dropout": transformers_config["hidden_dropout_prob"],
+        "position_embedding_type": transformers_config[
+            "position_embedding_type"
+        ],
+        "pad_token_id": transformers_config["pad_token_id"],
+        "max_sequence_length": transformers_config.get(
+            "max_position_embeddings", None
+        ),
+        "layer_norm_eps": transformers_config.get("layer_norm_eps", 1e-12),
+        "use_pre_layer_norm": transformers_config.get(
+            "emb_layer_norm_before", False
+        ),
+        "activation": transformers_config.get("activation", "gelu"),
+        "max_wavelength": transformers_config.get("max_wavelength", 10000),
+    }
+
+
+def transpose_and_reshape(x, shape):
+    return np.reshape(np.transpose(x), shape)
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embedding layer
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").embeddings,
+        hf_weight_key="embeddings.word_embeddings.weight",
+    )
+    if transformers_config["position_embedding_type"] == "absolute":
+        loader.port_weight(
+            keras_variable=backbone.get_layer(
+                "position_embedding"
+            ).position_embeddings,
+            hf_weight_key="embeddings.position_embeddings.weight",
+        )
+    if transformers_config.get("emb_layer_norm_before", False):
+        loader.port_weight(
+            keras_variable=backbone.get_layer("emb_layer_norm").gamma,
+            hf_weight_key="embeddings.layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=backbone.get_layer("emb_layer_norm").beta,
+            hf_weight_key="embeddings.layer_norm.bias",
+        )
+
+    loader.port_weight(
+        keras_variable=backbone.output_layer_norm.gamma,
+        hf_weight_key="encoder.emb_layer_norm_after.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.output_layer_norm.beta,
+        hf_weight_key="encoder.emb_layer_norm_after.bias",
+    )
+
+    # Attention blocks
+    for i in range(backbone.num_layers):
+        block = backbone.get_layer(f"transformer_layer_{i}")
+        attn = block.attention_layer
+        hf_prefix = "encoder.layer."
+        # Attention layers
+        loader.port_weight(
+            keras_variable=attn.q_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.query.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.q_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.query.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.k_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.key.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.k_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.key.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.v_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.value.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.v_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.value.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.o_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.dense.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.o_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.dense.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        # Attention layer norm.
+        loader.port_weight(
+            keras_variable=block.attention_norm.gamma,
+            hf_weight_key=f"{hf_prefix}{i}.attention.LayerNorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=block.attention_norm.beta,
+            hf_weight_key=f"{hf_prefix}{i}.attention.LayerNorm.bias",
+        )
+        # MLP layers
+        loader.port_weight(
+            keras_variable=block.feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.intermediate.dense.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_intermediate_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.intermediate.dense.bias",
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.output.dense.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_output_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.output.dense.bias",
+        )
+        # Output layer norm.
+        loader.port_weight(
+            keras_variable=block.feedforward_norm.gamma,
+            hf_weight_key=f"{hf_prefix}{i}.LayerNorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_norm.beta,
+            hf_weight_key=f"{hf_prefix}{i}.LayerNorm.bias",
+        )
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    return cls(
+        get_file(preset, "vocab.txt"),
+        lowercase=False,
+        **kwargs,
+    )

keras_hub/src/utils/transformers/convert_llama3.py

@@ -127,6 +127,12 @@ def convert_tokenizer(cls, preset, **kwargs):
     vocab = tokenizer_config["model"]["vocab"]
     merges = tokenizer_config["model"]["merges"]
 
+    # Handle different merge formats
+    if merges and isinstance(merges[0], list) and len(merges[0]) == 2:
+        # Convert list of lists format [["Ġ", "a"], ["Ġ", "b"]]
+        # to space-separated strings
+        merges = [" ".join(merge) for merge in merges]
+
     # Load all special tokens with the exception of "reserved" ones.
     special_tokens = set()
     for token in tokenizer_config["added_tokens"]:

keras_hub/src/utils/transformers/convert_qwen3.py

@@ -0,0 +1,145 @@
+import numpy as np
+
+from keras_hub.src.models.qwen3.qwen3_backbone import Qwen3Backbone
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = Qwen3Backbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "head_dim": transformers_config["head_dim"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "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"],
+        "layer_norm_epsilon": transformers_config["rms_norm_eps"],
+        "rope_max_wavelength": transformers_config["rope_theta"],
+        "sliding_window_size": transformers_config["sliding_window"]
+        if transformers_config["use_sliding_window"]
+        else None,
+        "tie_word_embeddings": transformers_config["tie_word_embeddings"],
+    }
+
+
+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
+        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)