Upload modeling.py with huggingface_hub

modeling.py

@@ -236,6 +236,17 @@ class GroupedQueryAttention(nn.Module):
         scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
 
         if attention_mask is not None:
+            # Handle different attention mask formats
+            if attention_mask.dim() == 4:  # Standard transformers format [batch, 1, 1, seq_len]
+                attention_mask = attention_mask.squeeze(1).squeeze(1)  # [batch, seq_len]
+                # Expand to match scores shape [batch, num_heads, seq_len, seq_len]
+                attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+                attention_mask = attention_mask.expand(-1, self.num_heads, seq_len, -1)
+            elif attention_mask.dim() == 2:  # [seq_len, seq_len] causal mask
+                # Expand to match scores shape [batch, num_heads, seq_len, seq_len]
+                attention_mask = attention_mask.unsqueeze(0).unsqueeze(0)
+                attention_mask = attention_mask.expand(batch_size, self.num_heads, -1, -1)
+
             scores = scores + attention_mask
 
         attention = F.softmax(scores, dim=-1)