pass attention mask

Author: DavidLandup0

keras_hub/src/models/smollm3/smollm3_backbone.py

@@ -90,7 +90,7 @@ def __init__(
                 layer_idx=i,
                 intermediate_size=intermediate_dim,
                 mlp_bias=mlp_bias,
-                rms_norm_epsilon=layer_norm_epsilon,
+                layer_norm_epsilon=layer_norm_epsilon,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)

keras_hub/src/models/smollm3/smollm3_layers.py

@@ -8,6 +8,7 @@
 from keras_hub.src.models.smollm3.smollm3_utils import rope_init
 
 
+
 class SmolLM3Attention(layers.Layer):
     """
     Multi-head attention layer for SmolLM3 model.
@@ -94,14 +95,14 @@ def build(self, input_shape):
         self.k_proj.build(hidden_states_shape)
         self.v_proj.build(hidden_states_shape)
         self.o_proj.build(hidden_states_shape)
-        self.training = False
         super().build(input_shape)
 
     def call(
         self,
         hidden_states,
         position_embeddings,
         training=False,
+        attention_mask=None,
         **kwargs,
     ):
         """
@@ -142,19 +143,14 @@ def _compute_kv_values(x_input):
             value_states = ops.transpose(value_states_raw, axes=(0, 2, 1, 3))
             return key_states, value_states
 
-        print("self_attention_cache is ", self_attention_cache)
         if self_attention_cache is not None:
             key_cache = self_attention_cache[:, 0, ...]
             value_cache = self_attention_cache[:, 1, ...]
 
             if self_attention_cache_update_index is None:
-                print("self_attention_cache_update_index is None")
                 key_states = key_cache
                 value_states = value_cache
             else:
-                print(
-                    "self_attention_cache_update_index is not None, computing kv values"
-                )
                 key_update, value_update = _compute_kv_values(hidden_states)
                 update_idx_tensor = ops.convert_to_tensor(
                     self_attention_cache_update_index, dtype="int32"
@@ -190,6 +186,7 @@ def _compute_kv_values(x_input):
             dropout=self.attention_dropout,
             scaling=self.scaling,
             training=self.training,
+            attention_mask=attention_mask,
         )
 
         attn_output = ops.reshape(attn_output, (*input_shape, self.hidden_size))
@@ -277,7 +274,6 @@ def build(self, input_shape):
             self.intermediate_size,
         )
         self.down_proj.build(down_proj_input_shape)
-        self.training = False
         super().build(input_shape)
 
     def call(self, x):
@@ -322,7 +318,7 @@ class SmolLM3DecoderLayer(layers.Layer):
         layer_idx: Index of the current layer.
         intermediate_size: The intermediate size of the MLP.
         mlp_bias: Whether to use bias in MLP dense layers.
-        rms_norm_epsilon: Epsilon for RMSNormalization.
+        layer_norm_epsilon: Epsilon for RMSNormalization.
     """
 
     def __init__(
@@ -337,7 +333,7 @@ def __init__(
         layer_idx: int,
         intermediate_size: int,
         mlp_bias: bool,
-        rms_norm_epsilon: float,
+        layer_norm_epsilon: float,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -364,10 +360,10 @@ def __init__(
         )
 
         self.input_layernorm = layers.RMSNormalization(
-            epsilon=rms_norm_epsilon, axis=-1, name="input_layernorm"
+            epsilon=layer_norm_epsilon, axis=-1, name="input_layernorm"
         )
         self.post_attention_layernorm = layers.RMSNormalization(
-            epsilon=rms_norm_epsilon, axis=-1, name="post_attention_layernorm"
+            epsilon=layer_norm_epsilon, axis=-1, name="post_attention_layernorm"
         )
 
         self.attention_type = layer_types[layer_idx]
@@ -399,7 +395,6 @@ def build(self, input_shape):
         self.mlp.build(input_shape)
         self.input_layernorm.build(input_shape)
         self.post_attention_layernorm.build(input_shape)
-        self.training = False
 
         super().build(input_shape)
 
@@ -518,7 +513,6 @@ def build(self, input_shape):
                          - position_ids_shape: (batch_size, seq_len)
         """
         # No internal layers to explicitly build here, as inv_freq is added in __init__
-        self.training = False
         super().build(input_shape)
 
     def call(

keras_hub/src/models/smollm3/smollm3_utils.py

@@ -38,7 +38,6 @@ def eager_attention_forward(
     dropout=0.0,
     training=False,
 ):
-    print('training', training)
     key_states = repeat_kv(key, module.num_key_value_groups)
     value_states = repeat_kv(value, module.num_key_value_groups)