Update modeling_quiet.py

modeling_quiet.py

@@ -54,21 +54,61 @@ _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:
-        if attention_mask.dim() == 3:
-            # Expanding from [batch_size, 1, tgt_len] to [batch_size, 1, tgt_len, tgt_len]
-            attention_mask = attention_mask.expand(bsz, 1, tgt_len, tgt_len)
+        # 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:
-            # Expanding from [batch_size, tgt_len] to [batch_size, 1, tgt_len, tgt_len]
-            attention_mask = attention_mask.unsqueeze(1).expand(bsz, 1, tgt_len, tgt_len)
+            # 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(f"Unexpected attention mask shape: {attention_mask.shape}, expected 2 or 3 dimensions.")
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
 
-        attention_mask = (1.0 - attention_mask) * -10000.0  # Masking operation for softmax
+    # 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 attention_mask
+    return combined_attention_mask
 
 
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data
@@ -1056,58 +1096,36 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         # Apply the language model head to get the final logits
         logits = self.lm_head(mixed_hidden_states)
         return logits
-
+    
     @torch.no_grad()
     def generate(
         self,
-        input_ids=None,
-        attention_mask=None,
-        max_new_tokens=None,
-        min_length=None,
-        do_sample=None,
-        early_stopping=None,
-        num_beams=None,
-        temperature=1.0,
-        top_k=None,
-        top_p=None,
-        repetition_penalty=None,
-        bad_words_ids=None,
-        bos_token_id=None,
-        pad_token_id=None,
-        eos_token_id=None,
-        length_penalty=None,
-        no_repeat_ngram_size=None,
-        num_return_sequences=None,
-        decoder_start_token_id=None,
-        use_cache=None,
-        num_beam_groups=None,
-        diversity_penalty=None,
-        prefix_allowed_tokens_fn=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        output_scores=None,
-        return_dict_in_generate=None,
-        forced_bos_token_id=None,
-        forced_eos_token_id=None,
-        remove_invalid_values=None,
-        synced_gpus=None,
-        **model_kwargs,
+        input_ids: torch.LongTensor = torch.LongTensor(),
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: Optional[int] = None,
+        temperature: float = 1.1,
+        **kwargs,
     ):
-        # Prepare the generation process with customized settings
-        model_inputs = self.prepare_inputs_for_generation(
-            input_ids, past_key_values=None, attention_mask=attention_mask, use_cache=use_cache, **model_kwargs
-        )
-        
-        from .generate import custom_generate
-        return custom_generate(
+        if isinstance(input_ids, str):
+            input_ids = self.tokenizer(input_ids, return_tensors="pt").input_ids
+
+        if attention_mask is None:
+            # Create a default attention mask if not provided
+            attention_mask = torch.ones_like(input_ids)
+
+        from .generate import generate
+
+        output = generate(
             self,
-            input_ids=input_ids,
+            input_ids,
             attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
-            **model_kwargs
+            **kwargs,
         )
 
+        return output.sequences
+
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(