Update app.py

app.py

@@ -29,6 +29,15 @@ st.markdown("""
         border-radius: 10px;
         padding: 15px;
     }
+    .architecture {
+        font-family: monospace;
+        color: #00ff00;
+        white-space: pre-wrap;
+        background-color: #1a1a1a;
+        padding: 20px;
+        border-radius: 10px;
+        border: 1px solid #00ff00;
+    }
 </style>
 """, unsafe_allow_html=True)
 
@@ -57,7 +66,6 @@ def plot_model_comparison(selected_model):
     fig, ax = plt.subplots(figsize=(10, 6))
     bars = ax.bar(model_names, params)
     
-    # Highlight selected model
     index = list(MODELS.keys()).index(selected_model)
     bars[index].set_color('#00ff00')
     
@@ -70,8 +78,70 @@ def plot_model_comparison(selected_model):
     
     st.pyplot(fig)
 
+def visualize_architecture(model_info):
+    architecture = []
+    model_type = model_info["type"]
+    layers = model_info["layers"]
+    heads = model_info["heads"]
+    
+    architecture.append("Input")
+    architecture.append("│")
+    architecture.append("▼")
+    
+    if model_type == "Encoder":
+        architecture.append("[Embedding Layer]")
+        for i in range(layers):
+            architecture.extend([
+                f"Encoder Layer {i+1}",
+                "├─ Multi-Head Attention",
+                f"│  └─ {heads} Heads",
+                "├─ Layer Normalization",
+                "└─ Feed Forward Network",
+                "│",
+                "▼"
+            ])
+        architecture.append("[Output]")
+    
+    elif model_type == "Decoder":
+        architecture.append("[Embedding Layer]")
+        for i in range(layers):
+            architecture.extend([
+                f"Decoder Layer {i+1}",
+                "├─ Masked Multi-Head Attention",
+                f"│  └─ {heads} Heads",
+                "├─ Layer Normalization",
+                "└─ Feed Forward Network",
+                "│",
+                "▼"
+            ])
+        architecture.append("[Output]")
+    
+    elif model_type == "Seq2Seq":
+        architecture.append("Encoder Stack")
+        for i in range(layers):
+            architecture.extend([
+                f"Encoder Layer {i+1}",
+                "├─ Self-Attention",
+                "└─ Feed Forward Network",
+                "│",
+                "▼"
+            ])
+        architecture.append("→→→ [Context] →→→")
+        architecture.append("Decoder Stack")
+        for i in range(layers):
+            architecture.extend([
+                f"Decoder Layer {i+1}",
+                "├─ Masked Self-Attention",
+                "├─ Encoder-Decoder Attention",
+                "└─ Feed Forward Network",
+                "│",
+                "▼"
+            ])
+        architecture.append("[Output]")
+    
+    return "\n".join(architecture)
+
 def visualize_attention_patterns():
-    # Simplified attention patterns visualization
     fig, ax = plt.subplots(figsize=(8, 6))
     data = torch.randn(5, 5)
     ax.imshow(data, cmap='viridis')
@@ -83,14 +153,10 @@ def visualize_attention_patterns():
 def main():
     st.title("🧠 Transformer Model Visualizer")
     
-    # Model selection
     selected_model = st.sidebar.selectbox("Select Model", list(MODELS.keys()))
-    
-    # Model details
     model_info = MODELS[selected_model]
     config = get_model_config(selected_model)
     
-    # Display metrics
     col1, col2, col3, col4 = st.columns(4)
     with col1:
         st.metric("Model Type", model_info["type"])
@@ -101,13 +167,20 @@ def main():
     with col4:
         st.metric("Parameters", f"{model_info['params']}M")
     
-    # Visualization tabs
-    tab1, tab2, tab3 = st.tabs(["Model Structure", "Comparison", "Model Specific"])
+    tab1, tab2, tab3 = st.tabs(["Model Structure", "Comparison", "Model Attention"])
     
     with tab1:
         st.subheader("Architecture Diagram")
-        st.image("https://jalammar.github.io/images/t/transformer.png",
-                  use_container_width=True)
+        architecture = visualize_architecture(model_info)
+        st.markdown(f"<div class='architecture'>{architecture}</div>", unsafe_allow_html=True)
+        
+        st.markdown("""
+        **Legend:**
+        - **Multi-Head Attention**: Self-attention mechanism with multiple parallel heads
+        - **Layer Normalization**: Normalization operation between layers
+        - **Feed Forward Network**: Position-wise fully connected network
+        - **Masked Attention**: Attention with future token masking
+        """)
     
     with tab2:
         st.subheader("Model Size Comparison")