Update modeling_quiet.py

modeling_quiet.py

@@ -54,61 +54,21 @@ _CONFIG_FOR_DOC = "QuietConfig"
 
 
 def _prepare_4d_causal_attention_mask_for_sdpa(attention_mask, input_shape, inputs_embeds, past_key_values_length):
-    # Compute the attention mask correctly
     bsz, tgt_len = input_shape
 
-    # Create a 4D attention mask from a 2D tensor mask.
-    # The shape of the output attention mask is (batch_size, 1, tgt_len, src_len)
-    # The values are either 0 or 1, where 0 means padding and 1 means non-padding.
-    combined_attention_mask = None
     if attention_mask is not None:
-        # What if attention_mask is not None and has a shape of (batch_size, 1, tgt_len, src_len)
-        # In this case, we can just use it directly.
-        if attention_mask.dim() == 4:
-            combined_attention_mask = attention_mask
-        # What if attention_mask is not None and has a shape of (batch_size, 1, tgt_len)
-        # In this case, we need to expand it to (batch_size, 1, tgt_len, src_len)
-        elif attention_mask.dim() == 3:
-            expanded_attn_mask = attention_mask[:, None, :, :]
-            combined_attention_mask = expanded_attn_mask
-        # What if attention_mask is not None and has a shape of (batch_size, tgt_len)
-        # In this case, we need to expand it to (batch_size, 1, tgt_len, src_len)
+        if attention_mask.dim() == 3:
+            # Expanding from [batch_size, 1, tgt_len] to [batch_size, 1, tgt_len, tgt_len]
+            attention_mask = attention_mask.expand(bsz, 1, tgt_len, tgt_len)
         elif attention_mask.dim() == 2:
-            # Provided a padding mask of dimensions [batch_size, seq_length]
-            # - if the model is a decoder, apply a causal mask in addition to the padding mask
-            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
-            if past_key_values_length > 0:
-                attention_mask = attention_mask.to(dtype=torch.long)
-                attention_mask = attention_mask[:, past_key_values_length:]
-            expanded_attn_mask = attention_mask[:, None, None, :]
-            combined_attention_mask = expanded_attn_mask
+            # Expanding from [batch_size, tgt_len] to [batch_size, 1, tgt_len, tgt_len]
+            attention_mask = attention_mask.unsqueeze(1).expand(bsz, 1, tgt_len, tgt_len)
         else:
-            raise ValueError(
-                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
-                    input_shape, attention_mask.shape
-                )
-            )
+            raise ValueError(f"Unexpected attention mask shape: {attention_mask.shape}, expected 2 or 3 dimensions.")
 
-    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
-    # masked positions, this operation will create a tensor which is 0.0 for
-    # positions we want to attend and -10000.0 for masked positions.
-    # Since we are adding it to the raw scores before the softmax, this is
-    # effectively the same as removing these entirely.
-    if combined_attention_mask is not None:
-        # Ensure the attention mask values are within a reasonable range
-        combined_attention_mask = combined_attention_mask.clamp(min=0, max=1)
-
-        # Convert the attention mask to bfloat16
-        combined_attention_mask = combined_attention_mask.to(torch.bfloat16)
-
-        # Normalize the attention mask values to be between 0 and 1
-        combined_attention_mask = (1.0 - combined_attention_mask) * -10000.0
-    else:
-        combined_attention_mask = torch.zeros(
-            (bsz, 1, tgt_len, tgt_len), dtype=torch.bfloat16, device=inputs_embeds.device
-        )
+        attention_mask = (1.0 - attention_mask) * -10000.0  # Masking operation for softmax
 
-    return combined_attention_mask
+    return attention_mask
 
 
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data