Update app.py

app.py

@@ -38,12 +38,13 @@ def estimate_transformer_stats(batch_size, seq_len, num_layers, hidden_dim, voca
     ]
 
     if show_breakdown:
+        total_all = attn_proj_flops + attn_score_flops + attn_out_proj_flops + ffn_flops + logit_flops
         output_lines.append("\nComponent-wise totals across training batch:")
-        output_lines.append(f"  - QKV Projections: {attn_proj_flops * B * L:.2e}")
-        output_lines.append(f"  - Attention Scores: {attn_score_flops * B * L:.2e}")
-        output_lines.append(f"  - Attention Output: {attn_out_proj_flops * B * L:.2e}")
-        output_lines.append(f"  - FFN: {ffn_flops * B * L:.2e}")
-        output_lines.append(f"  - Logits: {logit_flops * B * L:.2e}")
+        output_lines.append(f"  - QKV Projections: {attn_proj_flops * B * L:.2e} ({100 * attn_proj_flops / total_all:.1f}%)")
+        output_lines.append(f"  - Attention Scores: {attn_score_flops * B * L:.2e} ({100 * attn_score_flops / total_all:.1f}%)")
+        output_lines.append(f"  - Attention Output: {attn_out_proj_flops * B * L:.2e} ({100 * attn_out_proj_flops / total_all:.1f}%)")
+        output_lines.append(f"  - FFN: {ffn_flops * B * L:.2e} ({100 * ffn_flops / total_all:.1f}%)")
+        output_lines.append(f"  - Logits: {logit_flops * B * L:.2e} ({100 * logit_flops / total_all:.1f}%)")
 
     return "\n".join(output_lines)