Update modeling_quiet.py

modeling_quiet.py

@@ -1281,7 +1281,27 @@ class QuietForCausalLM(QuietPreTrainedModel):
         # Generate the continuation
         continuation_length = self.n_ahead - 2
         new_key_values = past_key_values
-        
+
+        if self.n_ahead != 1 or self.n_ahead_talk != 1 or self.comparison_mode:
+            if attention_mask is None:
+                base_attention_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=0).to(input_ids.device)
+                base_attention_mask = base_attention_mask.view(1, 1, seq_len, seq_len)
+                base_attention_mask = base_attention_mask.repeat(input_ids.shape[0], 1, 1, 1)
+                attention_mask = base_attention_mask
+            elif attention_mask.dim() == 2:
+                if seq_len + past_key_values_length != attention_mask.shape[-1]:
+                    attention_mask = torch.cat(
+                        [torch.ones((attention_mask.shape[0], past_key_values_length), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask],
+                        dim=-1
+                    )
+                attention_mask = _prepare_4d_causal_attention_mask(
+                    attention_mask,
+                    (batch_size, seq_len),
+                    inputs_embeds,
+                    past_key_values_length,
+                    sliding_window=self.config.sliding_window,
+                )
+    
         start_time = time.time()
         for continuation_idx in range(continuation_length):
             outputs = self.model(