Fix conversion weight names

Author: DavidLandup0

keras_hub/src/models/smollm3/smollm3_utils.py

@@ -48,7 +48,6 @@ def eager_attention_forward(
 
     # Apply attention mask if provided
     if attention_mask is not None:
-        # causal_mask = attention_mask[:, :, :, : ops.shape(key_states)[-2]]
         attn_weights = ops.add(attn_weights, attention_mask)
 
     attn_weights = ops.softmax(attn_weights, axis=-1)

keras_hub/src/utils/transformers/convert_smollm3.py

@@ -30,9 +30,6 @@ def convert_backbone_config(transformers_config):
         "rope_layer_enabled_list": transformers_config["no_rope_layers"],
         "layer_types": transformers_config["layer_types"],
         "mlp_bias": transformers_config["mlp_bias"],
-        "num_hidden_layers": transformers_config[
-            "num_hidden_layers"
-        ],  # Redundant with num_layers, but kept for completeness
     }
 
 
@@ -50,41 +47,41 @@ def transpose_and_reshape(x, shape):
 
         # Input layernorm
         loader.port_weight(
-            keras_variable=decoder_layer._self_attention_layernorm.scale,
+            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_attention_layer._query_dense.kernel,
+            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,
         )
         loader.port_weight(
-            keras_variable=decoder_layer._self_attention_layer._query_dense_layer_norm.scale,
+            keras_variable=decoder_layer.self_attn.q_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,
+            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,
         )
         loader.port_weight(
-            keras_variable=decoder_layer._self_attention_layer._key_dense_layer_norm.scale,
+            keras_variable=decoder_layer.self_attn.k_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,
+            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_attention_layer._output_dense.kernel,
+            keras_variable=decoder_layer.self_attn.o_proj.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},
@@ -93,27 +90,27 @@ def transpose_and_reshape(x, shape):
 
         # MLP layers
         loader.port_weight(
-            keras_variable=decoder_layer._feedforward_intermediate_dense.kernel,
+            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._feedforward_output_dense.kernel,
+            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._feedforward_gate_dense.kernel,
+            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._feedforward_layernorm.scale,
+            keras_variable=decoder_layer.post_attention_layernorm.scale,
             hf_weight_key=f"model.layers.{i}.post_attention_layernorm.weight",
         )