Update modeling_quiet.py

modeling_quiet.py

@@ -36,7 +36,7 @@ from transformers import TextStreamer, AutoTokenizer
 
 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 (
@@ -65,62 +65,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