Update app.py

app.py

@@ -75,30 +75,44 @@ class LlamaAttention(nn.Module):
     def forward(self, hidden_states, attention_mask=None):
         batch_size, seq_length, _ = hidden_states.size()
         
+        # Project and reshape
         q = self.q_proj(hidden_states).view(batch_size, seq_length, self.num_heads, self.head_dim)
         k = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
         v = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
         
+        # Repeat k/v heads if needed
         if self.num_kv_heads < self.num_heads:
             k = k.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
             v = v.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
         
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
+        # Transpose for attention
+        q = q.transpose(1, 2)  # (batch, num_heads, seq_len, head_dim)
+        k = k.transpose(1, 2)  # (batch, num_heads, seq_len, head_dim)
+        v = v.transpose(1, 2)  # (batch, num_heads, seq_len, head_dim)
         
-        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        # Calculate attention scores
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale  # (batch, num_heads, seq_len, seq_len)
         
+        # Apply attention mask if provided
         if attention_mask is not None:
-            attention_scores = attention_scores + attention_mask
-            
-        attention_probs = F.softmax(attention_scores, dim=-1)
-        context = torch.matmul(attention_probs, v)
+            # Ensure mask is broadcastable
+            if attention_mask.dim() == 2:
+                attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)  # (batch, 1, 1, seq_len)
+            scores = scores + attention_mask
+        
+        # Apply softmax and dropout
+        attention_weights = F.softmax(scores, dim=-1)
+        
+        # Apply attention to values
+        output = torch.matmul(attention_weights, v)  # (batch, num_heads, seq_len, head_dim)
         
-        context = context.transpose(1, 2).contiguous()
-        context = context.view(batch_size, seq_length, -1)
+        # Reshape and project back
+        output = output.transpose(1, 2).contiguous()  # (batch, seq_len, num_heads, head_dim)
+        output = output.view(batch_size, seq_length, -1)  # (batch, seq_len, hidden_size)
+        output = self.o_proj(output)
         
-        return self.o_proj(context)
+        return output
 
 class LlamaMLP(nn.Module):
     def __init__(self, config):