keras-hub-nightly 0.22.0.dev202508170419__py3-none-any.whl → 0.24.0.dev202511090424__py3-none-any.whl

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)

keras_hub/src/utils/transformers/convert_smollm3.py

@@ -0,0 +1,139 @@
+import numpy as np
+
+from keras_hub.src.models.smollm3.smollm3_backbone import SmolLM3Backbone
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = SmolLM3Backbone
+
+
+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_attention_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"
+        ],  # Using rms_norm_eps as layer_norm_epsilon
+        "max_position_embeddings": transformers_config[
+            "max_position_embeddings"
+        ],
+        "rope_theta": transformers_config["rope_theta"],
+        # partial_rotary_factor is not explicitly in config.json
+        # but is inherited from the default value in the
+        # `_compute_default_rope_parameters()` function
+        "partial_rotary_factor": 1.0,
+        "attention_bias": transformers_config["attention_bias"],
+        "attention_dropout": transformers_config["attention_dropout"],
+        # Despite the name, no_rope_layers: 1 = HAS RoPE, 0 = NO RoPE
+        "rope_layer_enabled_list": [
+            bool(x) for x in transformers_config["no_rope_layers"]
+        ],
+        "layer_types": transformers_config["layer_types"],
+        "mlp_bias": transformers_config["mlp_bias"],
+    }
+
+
+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",
+    )
+
+    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.input_layernorm.scale,
+            hf_weight_key=f"model.layers.{i}.input_layernorm.weight",
+        )
+
+        # Attention layers
+        ## Query
+        loader.port_weight(
+            keras_variable=decoder_layer.self_attn.q_proj.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_attn.k_proj.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_attn.v_proj.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_attn.o_proj.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.o_proj.weight",
+            hook_fn=transpose_and_reshape,
+        )
+
+        # MLP layers
+        loader.port_weight(
+            keras_variable=decoder_layer.mlp.up_proj.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.mlp.down_proj.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.mlp.gate_proj.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.post_attention_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",
+    )
+
+    backbone.training = False
+
+    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)

keras_hub/src/utils/transformers/convert_t5gemma.py

@@ -0,0 +1,229 @@
+from keras_hub.src.models.t5gemma.t5gemma_backbone import T5GemmaBackbone
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = T5GemmaBackbone
+
+
+def convert_backbone_config(transformers_config):
+    """Convert a Hugging Face T5Gemma config to a KerasHub backbone config."""
+    encoder_config = transformers_config["encoder"]
+    decoder_config = transformers_config["decoder"]
+
+    if decoder_config.get("hidden_activation") == "gelu_pytorch_tanh":
+        decoder_config["hidden_activation"] = "gelu_approximate"
+    if encoder_config.get("hidden_activation") == "gelu_pytorch_tanh":
+        encoder_config["hidden_activation"] = "gelu_approximate"
+
+    backbone_config = {
+        "vocabulary_size": decoder_config["vocab_size"],
+        "encoder_hidden_dim": encoder_config["hidden_size"],
+        "encoder_intermediate_dim": encoder_config["intermediate_size"],
+        "encoder_num_layers": encoder_config["num_hidden_layers"],
+        "encoder_num_attention_heads": encoder_config["num_attention_heads"],
+        "encoder_num_key_value_heads": encoder_config["num_key_value_heads"],
+        "encoder_head_dim": encoder_config["head_dim"],
+        "encoder_layer_types": encoder_config["layer_types"],
+        "decoder_hidden_dim": decoder_config["hidden_size"],
+        "decoder_intermediate_dim": decoder_config["intermediate_size"],
+        "decoder_num_layers": decoder_config["num_hidden_layers"],
+        "decoder_num_attention_heads": decoder_config["num_attention_heads"],
+        "decoder_num_key_value_heads": decoder_config["num_key_value_heads"],
+        "decoder_head_dim": decoder_config["head_dim"],
+        "decoder_layer_types": decoder_config["layer_types"],
+        "dropout_rate": decoder_config["dropout_rate"],
+        "rms_norm_eps": decoder_config["rms_norm_eps"],
+        "query_pre_attn_scalar": decoder_config["query_pre_attn_scalar"],
+        "tie_word_embeddings": transformers_config.get(
+            "tie_word_embeddings", True
+        ),
+        "attention_bias": decoder_config["attention_bias"],
+        "hidden_activation": decoder_config["hidden_activation"],
+        "initializer_range": decoder_config["initializer_range"],
+        "attention_dropout": decoder_config["attention_dropout"],
+        "sliding_window": decoder_config["sliding_window"],
+        "cross_attention_hidden_size": encoder_config["hidden_size"],
+        "attn_logit_softcapping": decoder_config["attn_logit_softcapping"],
+        "final_logit_softcapping": decoder_config["final_logit_softcapping"],
+        "rope_max_wavelength": decoder_config["rope_theta"],
+    }
+    return backbone_config
+
+
+def convert_weights(backbone, loader, transformers_config):
+    """Convert T5Gemma from Hugging Face to KerasHub."""
+    # Token embeddings.
+    loader.port_weight(
+        keras_variable=backbone.token_embedding.embeddings,
+        hf_weight_key="encoder.embed_tokens.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.decoder_token_embedding.embeddings,
+        hf_weight_key="decoder.embed_tokens.weight",
+    )
+
+    # Encoder.
+    loader.port_weight(
+        keras_variable=backbone.encoder_norm.scale,
+        hf_weight_key="encoder.norm.weight",
+    )
+    for i in range(backbone.encoder_num_layers):
+        layer = backbone.get_layer(f"encoder_layer_{i}")
+        hf_prefix = f"encoder.layers.{i}"
+
+        # Self-attention.
+        loader.port_weight(
+            keras_variable=layer.self_attn.query_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.q_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.self_attn.key_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.k_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.self_attn.value_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.v_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.self_attn.output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.o_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+
+        # MLP.
+        loader.port_weight(
+            keras_variable=layer.mlp.gate_proj.kernel,
+            hf_weight_key=f"{hf_prefix}.mlp.gate_proj.weight",
+            hook_fn=lambda w, s: w.T,
+        )
+        loader.port_weight(
+            keras_variable=layer.mlp.up_proj.kernel,
+            hf_weight_key=f"{hf_prefix}.mlp.up_proj.weight",
+            hook_fn=lambda w, s: w.T,
+        )
+        loader.port_weight(
+            keras_variable=layer.mlp.down_proj.kernel,
+            hf_weight_key=f"{hf_prefix}.mlp.down_proj.weight",
+            hook_fn=lambda w, s: w.T,
+        )
+
+        # Layer norm.
+        loader.port_weight(
+            keras_variable=layer.pre_self_attn_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.pre_self_attn_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.post_self_attn_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.post_self_attn_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.pre_feedforward_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.pre_feedforward_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.post_feedforward_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.post_feedforward_layernorm.weight",
+        )
+
+    # Decoder.
+    loader.port_weight(
+        keras_variable=backbone.decoder_norm.scale,
+        hf_weight_key="decoder.norm.weight",
+    )
+    for i in range(backbone.decoder_num_layers):
+        layer = backbone.get_layer(f"decoder_layer_{i}")
+        hf_prefix = f"decoder.layers.{i}"
+
+        # Self-attention.
+        loader.port_weight(
+            keras_variable=layer.self_attn.query_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.q_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.self_attn.key_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.k_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.self_attn.value_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.v_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.self_attn.output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.self_attn.o_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+
+        # Cross-attention.
+        loader.port_weight(
+            keras_variable=layer.cross_attn.query_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.cross_attn.q_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.cross_attn.key_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.cross_attn.k_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.cross_attn.value_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.cross_attn.v_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+        loader.port_weight(
+            keras_variable=layer.cross_attn.output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}.cross_attn.o_proj.weight",
+            hook_fn=lambda w, s: w.T.reshape(s),
+        )
+
+        # MLP.
+        loader.port_weight(
+            keras_variable=layer.mlp.gate_proj.kernel,
+            hf_weight_key=f"{hf_prefix}.mlp.gate_proj.weight",
+            hook_fn=lambda w, s: w.T,
+        )
+        loader.port_weight(
+            keras_variable=layer.mlp.up_proj.kernel,
+            hf_weight_key=f"{hf_prefix}.mlp.up_proj.weight",
+            hook_fn=lambda w, s: w.T,
+        )
+        loader.port_weight(
+            keras_variable=layer.mlp.down_proj.kernel,
+            hf_weight_key=f"{hf_prefix}.mlp.down_proj.weight",
+            hook_fn=lambda w, s: w.T,
+        )
+
+        # Layer norm.
+        loader.port_weight(
+            keras_variable=layer.pre_self_attn_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.pre_self_attn_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.post_self_attn_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.post_self_attn_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.pre_cross_attn_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.pre_cross_attn_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.post_cross_attn_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.post_cross_attn_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.pre_feedforward_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.pre_feedforward_layernorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=layer.post_feedforward_layernorm.scale,
+            hf_weight_key=f"{hf_prefix}.post_feedforward_layernorm.weight",
+        )
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    """Convert a T5Gemma tokenizer."""
+    return cls(get_file(preset, "tokenizer.model"), **kwargs)

keras_hub/src/utils/transformers/convert_vit.py

@@ -9,7 +9,10 @@ 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"],
+        "patch_size": (
+            transformers_config["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"],

keras_hub/src/utils/transformers/export/gemma.py

@@ -2,6 +2,7 @@ import keras.ops as ops
 
 
 def get_gemma_config(backbone):
+    token_embedding_layer = backbone.get_layer("token_embedding")
     hf_config = {
         "vocab_size": backbone.vocabulary_size,
         "num_hidden_layers": backbone.num_layers,
@@ -11,11 +12,16 @@ def get_gemma_config(backbone):
         "intermediate_size": backbone.intermediate_dim // 2,
         "head_dim": backbone.head_dim,
         "max_position_embeddings": 8192,
+        "tie_word_embeddings": token_embedding_layer.tie_weights,
+        "pad_token_id": 0,
+        "bos_token_id": 2,
+        "eos_token_id": 1,
+        "model_type": "gemma",
     }
     return hf_config
 
 
-def get_gemma_weights_map(backbone):
+def get_gemma_weights_map(backbone, include_lm_head=False):
     weights_dict = {}
 
     # Map token embedding
@@ -83,7 +89,46 @@ def get_gemma_weights_map(backbone):
         "final_normalization"
     ).weights[0]
 
-    # Tie weights, but clone to avoid sharing memory issues
-    weights_dict["lm_head.weight"] = ops.copy(token_embedding_layer.weights[0])
-
+    # Map lm_head if embeddings are not tied
+    if include_lm_head and not token_embedding_layer.tie_weights:
+        weights_dict["lm_head.weight"] = ops.transpose(
+            token_embedding_layer.reverse_embeddings
+        )
     return weights_dict
+
+
+def get_gemma_tokenizer_config(tokenizer):
+    tokenizer_config = {
+        "tokenizer_class": "GemmaTokenizer",
+        "clean_up_tokenization_spaces": False,
+        "bos_token": "<bos>",
+        "eos_token": "<eos>",
+        "pad_token": "<pad>",
+        "unk_token": "<unk>",
+        "add_bos_token": True,
+        "add_eos_token": False,
+        "model_max_length": 8192,
+    }
+    # Add added_tokens_decoder
+    added_tokens_decoder = {}
+    special_tokens = [
+        "<pad>",
+        "<bos>",
+        "<eos>",
+        "<unk>",
+        "<start_of_turn>",
+        "<end_of_turn>",
+    ]
+    for token in special_tokens:
+        token_id = tokenizer.token_to_id(token)
+        if token_id is not None:
+            added_tokens_decoder[str(token_id)] = {
+                "content": token,
+                "special": True,
+                "single_word": False,
+                "lstrip": False,
+                "rstrip": False,
+                "normalized": False,
+            }
+    tokenizer_config["added_tokens_decoder"] = added_tokens_decoder
+    return tokenizer_config