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, _prepare_4d_causal_attention_mask_for_sdpa
+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