keras-hub-nightly 0.21.0.dev202505040408__py3-none-any.whl → 0.21.0.dev202505060405__py3-none-any.whl

keras_hub/src/models/retinanet/retinanet_image_converter.py

@@ -6,16 +6,3 @@ from keras_hub.src.models.retinanet.retinanet_backbone import RetinaNetBackbone
 @keras_hub_export("keras_hub.layers.RetinaNetImageConverter")
 class RetinaNetImageConverter(ImageConverter):
     backbone_cls = RetinaNetBackbone
-
-    def __init__(
-        self,
-        *args,
-        **kwargs,
-    ):
-        # TODO: update presets and remove these old config options. They were
-        # never needed.
-        if "norm_mean" in kwargs:
-            kwargs["offset"] = [-x for x in kwargs.pop("norm_mean")]
-        if "norm_std" in kwargs:
-            kwargs["scale"] = [1.0 / x for x in kwargs.pop("norm_std")]
-        super().__init__(*args, **kwargs)

keras_hub/src/models/retinanet/retinanet_presets.py

@@ -11,7 +11,7 @@ backbone_presets = {
             "params": 34121239,
             "path": "retinanet",
         },
-        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_coco/3",
+        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_coco/4",
     },
     "retinanet_resnet50_fpn_v2_coco": {
         "metadata": {
@@ -22,6 +22,6 @@ backbone_presets = {
             "params": 31558592,
             "path": "retinanet",
         },
-        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_v2_coco/2",
+        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_v2_coco/3",
     },
 }

keras_hub/src/models/task.py

@@ -236,14 +236,17 @@ class Task(PipelineModel):
             objects_to_skip=backbone_layer_ids,
         )
 
-    def save_to_preset(self, preset_dir):
+    def save_to_preset(self, preset_dir, max_shard_size=10):
         """Save task to a preset directory.
 
         Args:
             preset_dir: The path to the local model preset directory.
+            max_shard_size: `int` or `float`. Maximum size in GB for each
+                sharded file. If `None`, no sharding will be done. Defaults to
+                `10`.
         """
         saver = get_preset_saver(preset_dir)
-        saver.save_task(self)
+        saver.save_task(self, max_shard_size=max_shard_size)
 
     @property
     def layers(self):

keras_hub/src/utils/keras_utils.py

@@ -1,3 +1,4 @@
+import inspect
 import sys
 
 import keras
@@ -147,3 +148,13 @@ def get_gpu_names():
         ]
     else:
         return [""]
+
+
+def sharded_weights_available():
+    """Whether sharded weights serialization is available.
+
+    Returns:
+        `True` if sharded weights are available, `False` otherwise.
+    """
+    save_weights_signature = inspect.signature(keras.saving.save_weights)
+    return "max_shard_size" in save_weights_signature.parameters

keras_hub/src/utils/preset_utils.py

@@ -10,6 +10,8 @@ from absl import logging
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.utils.keras_utils import print_msg
+from keras_hub.src.utils.keras_utils import sharded_weights_available
+from keras_hub.src.utils.tensor_utils import get_tensor_size_in_bits
 
 try:
     import kagglehub
@@ -48,6 +50,7 @@ METADATA_FILE = "metadata.json"
 # Weight file names.
 MODEL_WEIGHTS_FILE = "model.weights.h5"
 TASK_WEIGHTS_FILE = "task.weights.h5"
+SHARDED_MODEL_WEIGHTS_CONFIG_FILE = "model.weights.json"
 
 # HuggingFace filenames.
 README_FILE = "README.md"
@@ -647,7 +650,7 @@ class KerasPresetLoader(PresetLoader):
         backbone = self._load_serialized_object(self.config, **kwargs)
         if load_weights:
             jax_memory_cleanup(backbone)
-            backbone.load_weights(get_file(self.preset, MODEL_WEIGHTS_FILE))
+            self._load_backbone_weights(backbone)
         return backbone
 
     def load_tokenizer(self, cls, config_file=TOKENIZER_CONFIG_FILE, **kwargs):
@@ -697,8 +700,7 @@ class KerasPresetLoader(PresetLoader):
                 task.load_task_weights(task_weights)
             else:
                 jax_memory_cleanup(task.backbone)
-            backbone_weights = get_file(self.preset, MODEL_WEIGHTS_FILE)
-            task.backbone.load_weights(backbone_weights)
+            self._load_backbone_weights(task.backbone)
         return task
 
     def load_preprocessor(
@@ -726,18 +728,64 @@ class KerasPresetLoader(PresetLoader):
         config["config"] = {**config["config"], **kwargs}
         return keras.saving.deserialize_keras_object(config)
 
+    def _get_sharded_filenames(self, config_path):
+        with open(config_path, encoding="utf-8") as config_file:
+            config = json.load(config_file)
+        weight_map = config["weight_map"]
+        return sorted(set(weight_map.values()))
+
+    def _load_backbone_weights(self, backbone):
+        # Detect if the backbone is sharded or not.
+        has_single_file_weights = check_file_exists(
+            self.preset, MODEL_WEIGHTS_FILE
+        )
+        if has_single_file_weights:
+            filepath = get_file(self.preset, MODEL_WEIGHTS_FILE)
+        else:
+            if not sharded_weights_available():
+                raise RuntimeError(
+                    "Sharded weights loading is not supported in the current "
+                    f"Keras version {keras.__version__}. "
+                    "Please update to a newer version."
+                )
+            filepath = get_file(self.preset, SHARDED_MODEL_WEIGHTS_CONFIG_FILE)
+            sharded_filenames = self._get_sharded_filenames(filepath)
+            for sharded_filename in sharded_filenames:
+                # Download the sharded weights.
+                _ = get_file(self.preset, sharded_filename)
+        backbone.load_weights(filepath)
+
 
 class KerasPresetSaver:
     def __init__(self, preset_dir):
         os.makedirs(preset_dir, exist_ok=True)
         self.preset_dir = preset_dir
 
-    def save_backbone(self, backbone):
+    def save_backbone(self, backbone, max_shard_size=10):
         self._save_serialized_object(backbone, config_file=CONFIG_FILE)
-        backbone_weight_path = os.path.join(self.preset_dir, MODEL_WEIGHTS_FILE)
-        backbone.save_weights(backbone_weight_path)
         self._save_metadata(backbone)
 
+        # Save the weights.
+        backbone_size_in_bytes = self._get_variables_size_in_bytes(
+            backbone.variables
+        )
+        backbone_size_in_gb = backbone_size_in_bytes / (1024**3)
+        # If the size of the backbone is larger than `max_shard_size`, save
+        # sharded weights.
+        if sharded_weights_available() and backbone_size_in_gb > max_shard_size:
+            backbone_sharded_weights_config_path = os.path.join(
+                self.preset_dir, SHARDED_MODEL_WEIGHTS_CONFIG_FILE
+            )
+            backbone.save_weights(
+                backbone_sharded_weights_config_path,
+                max_shard_size=max_shard_size,
+            )
+        else:
+            backbone_weight_path = os.path.join(
+                self.preset_dir, MODEL_WEIGHTS_FILE
+            )
+            backbone.save_weights(backbone_weight_path)
+
     def save_tokenizer(self, tokenizer):
         config_file = TOKENIZER_CONFIG_FILE
         if hasattr(tokenizer, "config_file"):
@@ -755,7 +803,7 @@ class KerasPresetSaver:
     def save_image_converter(self, converter):
         self._save_serialized_object(converter, IMAGE_CONVERTER_CONFIG_FILE)
 
-    def save_task(self, task):
+    def save_task(self, task, max_shard_size=10):
         # Save task specific config and weights.
         self._save_serialized_object(task, TASK_CONFIG_FILE)
         if task.has_task_weights():
@@ -763,10 +811,12 @@ class KerasPresetSaver:
             task.save_task_weights(task_weight_path)
         # Save backbone.
         if hasattr(task.backbone, "save_to_preset"):
-            task.backbone.save_to_preset(self.preset_dir)
+            task.backbone.save_to_preset(
+                self.preset_dir, max_shard_size=max_shard_size
+            )
         else:
             # Allow saving a `keras.Model` that is not a backbone subclass.
-            self.save_backbone(task.backbone)
+            self.save_backbone(task.backbone, max_shard_size=max_shard_size)
         # Save preprocessor.
         if task.preprocessor and hasattr(task.preprocessor, "save_to_preset"):
             task.preprocessor.save_to_preset(self.preset_dir)
@@ -823,3 +873,13 @@ class KerasPresetSaver:
         metadata_path = os.path.join(self.preset_dir, METADATA_FILE)
         with open(metadata_path, "w") as metadata_file:
             metadata_file.write(json.dumps(metadata, indent=4))
+
+    def _get_variables_size_in_bytes(self, variables):
+        unique_variables = {}
+        for v in variables:
+            if id(v) not in unique_variables:
+                unique_variables[id(v)] = (v.shape, v.dtype)
+        total_memory_size = 0
+        for shape, dtype in unique_variables.values():
+            total_memory_size += get_tensor_size_in_bits(shape, dtype)
+        return total_memory_size / 8

keras_hub/src/utils/tensor_utils.py

@@ -1,6 +1,8 @@
 import contextlib
 import functools
 import inspect
+import math
+import re
 import threading
 
 import keras
@@ -305,6 +307,29 @@ def is_string_dtype(dtype):
     return "string" in keras.backend.standardize_dtype(dtype)
 
 
+def get_dtype_size_in_bits(dtype):
+    """Get the size of a given dtype in bits."""
+    dtype = keras.backend.standardize_dtype(dtype)
+    # If dtype is bool, return 1 immediately.
+    if dtype == "bool":
+        return 1
+    # Else, we extract the bit size from the string.
+    return int(re.sub(r"bfloat|float|uint|int", "", dtype))
+
+
+def get_tensor_size_in_bits(shape, dtype):
+    """Calculate the size given dtype and shape in bits.
+
+    Args:
+        dtype: The dtype of the tensor.
+        shape: List of iterables representing the shape of the tensor.
+
+    Returns:
+        The size of the tensor in bytes.
+    """
+    return math.prod(shape) * get_dtype_size_in_bits(dtype)
+
+
 def any_equal(inputs, values, padding_mask):
     """Return a mask that is True anywhere `inputs` has a value in `values`.
 
@@ -320,7 +345,8 @@ def any_equal(inputs, values, padding_mask):
     Returns:
         A tensor with `inputs` shape where each position is True if it contains
             a value from any `values`. Padding mask will be applied before
-            returning."""
+            returning.
+    """
     output = ops.equal(inputs, values[0])
     for value in values[1:]:
         value_equality = ops.equal(inputs, value)

keras_hub/src/utils/transformers/convert_mixtral.py

@@ -0,0 +1,139 @@
+import numpy as np
+
+from keras_hub.src.models.mixtral.mixtral_backbone import MixtralBackbone
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = MixtralBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_query_heads": transformers_config["num_attention_heads"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "num_key_value_heads": transformers_config["num_key_value_heads"],
+        "num_experts": transformers_config["num_local_experts"],
+        "top_k": transformers_config["num_experts_per_tok"],
+        "rope_max_wavelength": transformers_config["rope_theta"],
+        "layer_norm_epsilon": transformers_config["rms_norm_eps"],
+        "sliding_window": transformers_config["sliding_window"],
+        "output_router_logits": transformers_config["output_router_logits"],
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embeddings
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").embeddings,
+        hf_weight_key="model.embed_tokens.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").reverse_embeddings,
+        hf_weight_key="lm_head.weight",
+        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,
+        )
+        ## 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,
+        )
+        ## 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",
+            hook_fn=transpose_and_reshape,
+        )
+
+        # MoE layers
+        # Router gate
+        loader.port_weight(
+            keras_variable=decoder_layer._sparse_moe_block._sparse_feedforward_gate_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.block_sparse_moe.gate.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+
+        # Batched experts: w1 (gate), w3 (intermediate), and w2 (output) weights
+        gate_weights_list = []
+        intermediate_weights_list = []
+        output_weights_list = []
+        for expert_idx in range(backbone.num_experts):
+            # Load w1 (gate dense) for each expert
+            w1 = loader.get_tensor(
+                f"model.layers.{i}.block_sparse_moe.experts.{expert_idx}.w1.weight"
+            )
+            w1_transposed = np.transpose(w1, axes=(1, 0))
+            gate_weights_list.append(w1_transposed)
+
+            w3 = loader.get_tensor(
+                f"model.layers.{i}.block_sparse_moe.experts.{expert_idx}.w3.weight"
+            )
+            w3_transposed = np.transpose(w3, axes=(1, 0))
+            intermediate_weights_list.append(w3_transposed)
+
+            w2 = loader.get_tensor(
+                f"model.layers.{i}.block_sparse_moe.experts.{expert_idx}.w2.weight"
+            )
+            w2_transposed = np.transpose(w2, axes=(1, 0))
+            output_weights_list.append(w2_transposed)
+
+        gate_batched = np.stack(gate_weights_list, axis=0)
+        intermediate_batched = np.stack(intermediate_weights_list, axis=0)
+        output_batched = np.stack(output_weights_list, axis=0)
+
+        # Assign batched weights to expert_bank
+        decoder_layer._sparse_moe_block.expert_bank._expert_feedforward_gate_dense.assign(
+            gate_batched
+        )
+        decoder_layer._sparse_moe_block.expert_bank._expert_feedforward_intermediate_dense.assign(
+            intermediate_batched
+        )
+        decoder_layer._sparse_moe_block.expert_bank._expert_feedforward_output_dense.assign(
+            output_batched
+        )
+
+        # 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):
+    return cls(get_file(preset, "tokenizer.model"), **kwargs)

keras_hub/src/utils/transformers/convert_qwen_moe.py

@@ -0,0 +1,253 @@
+import numpy as np
+
+from keras_hub.src.models.qwen_moe.qwen_moe_backbone import QwenMoeBackbone
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = QwenMoeBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "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"],
+        "moe_intermediate_dim": transformers_config["moe_intermediate_size"],
+        "shared_expert_intermediate_dim": transformers_config[
+            "shared_expert_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"],
+        "use_sliding_window": transformers_config["use_sliding_window"],
+        "sliding_window_size": transformers_config["sliding_window"],
+        "output_router_logits": transformers_config["output_router_logits"],
+        "router_aux_loss_coefficient": transformers_config[
+            "router_aux_loss_coef"
+        ],
+    }
+
+
+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.bias,
+            hf_weight_key=f"model.layers.{i}.self_attn.q_proj.bias",
+            hook_fn=transpose_and_reshape,
+        )
+        ## 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.bias,
+            hf_weight_key=f"model.layers.{i}.self_attn.k_proj.bias",
+            hook_fn=transpose_and_reshape,
+        )
+        ## 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,
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._value_dense.bias,
+            hf_weight_key=f"model.layers.{i}.self_attn.v_proj.bias",
+            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
+            )
+
+            loader.port_weight(
+                keras_variable=decoder_layer.mlp.shared_expert_dense._feedforward_intermediate_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.shared_expert.up_proj.weight",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            loader.port_weight(
+                keras_variable=decoder_layer.mlp.shared_expert_dense._feedforward_output_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.shared_expert.down_proj.weight",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+            loader.port_weight(
+                keras_variable=decoder_layer.mlp.shared_expert_dense._feedforward_gate_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.shared_expert.gate_proj.weight",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+
+            loader.port_weight(
+                keras_variable=decoder_layer.mlp.shared_expert_gate_dense.kernel,
+                hf_weight_key=f"model.layers.{i}.mlp.shared_expert_gate.weight",
+                hook_fn=lambda hf_tensor, _: np.transpose(
+                    hf_tensor, axes=(1, 0)
+                ),
+            )
+        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"]
+
+    # 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)

keras_hub/src/utils/transformers/preset_loader.py

@@ -11,8 +11,10 @@ from keras_hub.src.utils.transformers import convert_gemma
 from keras_hub.src.utils.transformers import convert_gpt2
 from keras_hub.src.utils.transformers import convert_llama3
 from keras_hub.src.utils.transformers import convert_mistral
+from keras_hub.src.utils.transformers import convert_mixtral
 from keras_hub.src.utils.transformers import convert_pali_gemma
 from keras_hub.src.utils.transformers import convert_qwen
+from keras_hub.src.utils.transformers import convert_qwen_moe
 from keras_hub.src.utils.transformers import convert_vit
 from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
 
@@ -44,6 +46,10 @@ class TransformersPresetLoader(PresetLoader):
             self.converter = convert_vit
         elif model_type == "qwen2":
             self.converter = convert_qwen
+        elif model_type == "mixtral":
+            self.converter = convert_mixtral
+        elif model_type == "qwen2_moe":
+            self.converter = convert_qwen_moe
         else:
             raise ValueError(
                 "KerasHub has no converter for huggingface/transformers models "