Update modeling_quiet.py

modeling_quiet.py

@@ -37,7 +37,7 @@ import transformers
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
-from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import (
@@ -58,62 +58,62 @@ logger = logging.get_logger(__name__)
 _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)
-        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
-        else:
-            raise ValueError(
-                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
-                    input_shape, attention_mask.shape
-                )
-            )
-
-    # 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
-        )
-
-    return combined_attention_mask
+# 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)
+#         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
+#         else:
+#             raise ValueError(
+#                 "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+#                     input_shape, attention_mask.shape
+#                 )
+#             )
+
+#     # 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
+#         )
+
+#     return combined_attention_mask
 
 
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data