Update app.py

app.py

@@ -1,3 +1,10 @@
+from torch.nn.functional import softmax
+
+def simple_attention(query, key, value):
+    # Scaled dot-product attention
+    scores = query @ key.transpose(-2, -1) / (query.size(-1) ** 0.5)
+    attn_weights = softmax(scores, dim=-1)
+    return attn_weights @ value
 
 
 import spaces