Update modeling_quiet.py

modeling_quiet.py

@@ -1071,10 +1071,12 @@ class QuietModel(QuietPreTrainedModel):
                     " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
                 )
         if self._attn_implementation == "flash_attention_2":
+            
+            seq_length = original_sequence_length + num_thought_tokens
             # Convert 2D mask to 4D and adjust size
             attention_mask = _prepare_4d_causal_attention_mask(
                 attention_mask, 
-                (batch_size, seq_len + num_thought_tokens),  # Adjust size
+                (batch_size, seq_len),  # Adjust size
                 inputs_embeds, 
                 past_key_values_length, 
                 sliding_window=self.config.sliding_window,
@@ -1084,17 +1086,19 @@ class QuietModel(QuietPreTrainedModel):
         elif self._attn_implementation == "sdpa" and not output_attentions and attention_mask.dim() == 2 and False:
             # output_attentions=True can not be supported when using SDPA, and we fall back on
             # the manual implementation that requires a 4D causal mask in all cases.
+            seq_length = original_sequence_length + num_thought_tokens
             attention_mask = _prepare_4d_causal_attention_mask(
                 attention_mask, 
-                (batch_size, seq_len + num_thought_tokens),
+                (batch_size, seq_len),
                 inputs_embeds,
                 past_key_values_length,
             )
         elif attention_mask is None or attention_mask.dim() == 2:
             # 4d mask is passed through the layers
+            seq_length = original_sequence_length + num_thought_tokens
             attention_mask = _prepare_4d_causal_attention_mask(
                 attention_mask, 
-                (batch_size, seq_len + num_thought_tokens),
+                (batch_size, seq_len),
                 inputs_embeds,
                 past_key_values_length,
                 sliding_window=self.config.sliding_window,
@@ -1665,10 +1669,12 @@ class QuietForCausalLM(QuietPreTrainedModel):
                         num_thought_tokens += 1
                     if self.use_end_thought_token:
                         num_thought_tokens += 1
+                        
+                    seq_length = original_sequence_length + num_thought_tokens
                     # # if the attention mask 
                     attention_mask = _prepare_4d_causal_attention_mask(
                         attention_mask, 
-                        (batch_size, seq_len + num_thought_tokens),
+                        (batch_size, seq_len),
                         inputs_embeds,
                         past_key_values_length,
                         sliding_window=self.config.sliding_window,