Update app.py

app.py

@@ -2,9 +2,6 @@ import streamlit as st
 import matplotlib.pyplot as plt
 import pandas as pd
 import torch
-import plotly.express as px
-from sklearn.decomposition import PCA
-from sklearn.manifold import TSNE
 from transformers import AutoConfig, AutoTokenizer
 
 # Page configuration
@@ -49,38 +46,18 @@ st.markdown("""
 </style>
 """, unsafe_allow_html=True)
 
-# Enhanced Model database
+# Model database
 MODELS = {
-    "BERT": {"model_name": "bert-base-uncased", "type": "Encoder", "layers": 12, "heads": 12, 
-            "params": 109.48, "downloads": "10M+", "release_year": 2018, "gpu_req": "4GB+", 
-            "cpu_req": "4 cores+", "ram_req": "8GB+"},
-    "GPT-2": {"model_name": "gpt2", "type": "Decoder", "layers": 12, "heads": 12, 
-             "params": 117, "downloads": "8M+", "release_year": 2019, "gpu_req": "6GB+", 
-             "cpu_req": "4 cores+", "ram_req": "12GB+"},
-    "T5-Small": {"model_name": "t5-small", "type": "Seq2Seq", "layers": 6, "heads": 8, 
-                "params": 60, "downloads": "5M+", "release_year": 2019, "gpu_req": "3GB+", 
-                "cpu_req": "2 cores+", "ram_req": "6GB+"},
-    "RoBERTa": {"model_name": "roberta-base", "type": "Encoder", "layers": 12, "heads": 12, 
-               "params": 125, "downloads": "7M+", "release_year": 2019, "gpu_req": "5GB+", 
-               "cpu_req": "4 cores+", "ram_req": "10GB+"},
-    "DistilBERT": {"model_name": "distilbert-base-uncased", "type": "Encoder", "layers": 6, 
-                  "heads": 12, "params": 66, "downloads": "9M+", "release_year": 2019, 
-                  "gpu_req": "2GB+", "cpu_req": "2 cores+", "ram_req": "4GB+"},
-    "ALBERT": {"model_name": "albert-base-v2", "type": "Encoder", "layers": 12, "heads": 12, 
-              "params": 11.8, "downloads": "3M+", "release_year": 2019, "gpu_req": "1GB+", 
-              "cpu_req": "1 core+", "ram_req": "2GB+"},
-    "ELECTRA": {"model_name": "google/electra-small-discriminator", "type": "Encoder", 
-               "layers": 12, "heads": 12, "params": 13.5, "downloads": "2M+", 
-               "release_year": 2020, "gpu_req": "2GB+", "cpu_req": "2 cores+", "ram_req": "4GB+"},
-    "XLNet": {"model_name": "xlnet-base-cased", "type": "AutoRegressive", "layers": 12, 
-             "heads": 12, "params": 110, "downloads": "4M+", "release_year": 2019, 
-             "gpu_req": "5GB+", "cpu_req": "4 cores+", "ram_req": "8GB+"},
-    "BART": {"model_name": "facebook/bart-base", "type": "Seq2Seq", "layers": 6, "heads": 16, 
-            "params": 139, "downloads": "6M+", "release_year": 2020, "gpu_req": "6GB+", 
-            "cpu_req": "4 cores+", "ram_req": "12GB+"},
-    "DeBERTa": {"model_name": "microsoft/deberta-base", "type": "Encoder", "layers": 12, 
-               "heads": 12, "params": 139, "downloads": "3M+", "release_year": 2021, 
-               "gpu_req": "8GB+", "cpu_req": "6 cores+", "ram_req": "16GB+"}
+    "BERT": {"model_name": "bert-base-uncased", "type": "Encoder", "layers": 12, "heads": 12, "params": 109.48},
+    "GPT-2": {"model_name": "gpt2", "type": "Decoder", "layers": 12, "heads": 12, "params": 117},
+    "T5-Small": {"model_name": "t5-small", "type": "Seq2Seq", "layers": 6, "heads": 8, "params": 60},
+    "RoBERTa": {"model_name": "roberta-base", "type": "Encoder", "layers": 12, "heads": 12, "params": 125},
+    "DistilBERT": {"model_name": "distilbert-base-uncased", "type": "Encoder", "layers": 6, "heads": 12, "params": 66},
+    "ALBERT": {"model_name": "albert-base-v2", "type": "Encoder", "layers": 12, "heads": 12, "params": 11.8},
+    "ELECTRA": {"model_name": "google/electra-small-discriminator", "type": "Encoder", "layers": 12, "heads": 12, "params": 13.5},
+    "XLNet": {"model_name": "xlnet-base-cased", "type": "AutoRegressive", "layers": 12, "heads": 12, "params": 110},
+    "BART": {"model_name": "facebook/bart-base", "type": "Seq2Seq", "layers": 6, "heads": 16, "params": 139},
+    "DeBERTa": {"model_name": "microsoft/deberta-base", "type": "Encoder", "layers": 12, "heads": 12, "params": 139}
 }
 
 def get_model_config(model_name):
@@ -109,7 +86,7 @@ def plot_model_comparison(selected_model):
 def visualize_architecture(model_info):
     architecture = []
     model_type = model_info["type"]
-    layers = model_info["layers"]
+    layers = model_info.get("layers", model_info.get("layers", 12))  # Handle key variations
     heads = model_info["heads"]
     
     architecture.append("Input")
@@ -178,86 +155,13 @@ def visualize_attention_patterns():
     fig.patch.set_facecolor('#2c2c2c')
     st.pyplot(fig)
 
-def embedding_projector():
-    st.subheader("🔍 Embedding Projector")
-    
-    # Sample words for visualization
-    words = ["king", "queen", "man", "woman", "computer", "algorithm", 
-            "neural", "network", "language", "processing"]
-    
-    # Create dummy embeddings (3D for visualization)
-    embeddings = torch.randn(len(words), 256)
-    
-    # Dimensionality reduction
-    method = st.selectbox("Reduction Method", ["PCA", "t-SNE"])
-    
-    if method == "PCA":
-        reduced = PCA(n_components=3).fit_transform(embeddings)
+def get_hardware_recommendation(params):
+    if params < 100:
+        return "CPU or Entry-level GPU (e.g., GTX 1060)"
+    elif 100 <= params < 200:
+        return "Mid-range GPU (e.g., RTX 2080, RTX 3060)"
     else:
-        reduced = TSNE(n_components=3).fit_transform(embeddings.numpy())
-        
-    # Create interactive 3D plot
-    fig = px.scatter_3d(
-        x=reduced[:,0], y=reduced[:,1], z=reduced[:,2],
-        text=words,
-        title=f"Word Embeddings ({method})"
-    )
-    fig.update_traces(marker=dict(size=5), textposition='top center')
-    st.plotly_chart(fig, use_container_width=True)
-
-def hardware_recommendations(model_info):
-    st.subheader("💻 Hardware Recommendations")
-    
-    col1, col2, col3 = st.columns(3)
-    with col1:
-        st.metric("Minimum GPU", model_info.get("gpu_req", "4GB+"))
-    with col2:
-        st.metric("CPU Recommendation", model_info.get("cpu_req", "4 cores+"))
-    with col3:
-        st.metric("RAM Requirement", model_info.get("ram_req", "8GB+"))
-    
-    st.markdown("""
-    **Cloud Recommendations:**
-    - AWS: g4dn.xlarge instance
-    - GCP: n1-standard-4 with T4 GPU
-    - Azure: Standard_NC4as_T4_v3
-    """)
-
-def model_zoo_statistics():
-    st.subheader("📊 Model Zoo Statistics")
-    
-    df = pd.DataFrame.from_dict(MODELS, orient='index')
-    st.dataframe(
-        df[["release_year", "downloads", "params"]],
-        column_config={
-            "release_year": "Release Year",
-            "downloads": "Downloads",
-            "params": "Params (M)"
-        },
-        use_container_width=True,
-        height=400
-    )
-    
-    fig = px.bar(df, x=df.index, y="params", title="Model Parameters Comparison")
-    st.plotly_chart(fig, use_container_width=True)
-
-def memory_usage_estimator(model_info):
-    st.subheader("🧮 Memory Usage Estimator")
-    
-    precision = st.selectbox("Precision", ["FP32", "FP16", "INT8"])
-    batch_size = st.slider("Batch size", 1, 128, 8)
-    
-    # Memory calculation
-    bytes_map = {"FP32": 4, "FP16": 2, "INT8": 1}
-    estimated_memory = (model_info["params"] * 1e6 * bytes_map[precision] * batch_size) / (1024**3)
-    
-    col1, col2 = st.columns(2)
-    with col1:
-        st.metric("Estimated VRAM", f"{estimated_memory:.1f} GB")
-    with col2:
-        st.metric("Recommended GPU", "RTX 3090" if estimated_memory > 24 else "RTX 3060")
-    
-    st.progress(min(estimated_memory/40, 1.0), text="GPU Memory Utilization (of 40GB GPU)")
+        return "High-end GPU (e.g., RTX 3090, A100) or TPU"
 
 def main():
     st.title("🧠 Transformer Model Visualizer")
@@ -271,16 +175,15 @@ def main():
     with col1:
         st.metric("Model Type", model_info["type"])
     with col2:
-        st.metric("Layers", model_info["layers"])
+        st.metric("Layers", model_info.get("layers", model_info.get("layers", "N/A")))
     with col3:
         st.metric("Attention Heads", model_info["heads"])
     with col4:
         st.metric("Parameters", f"{model_info['params']}M")
     
-    # Updated tabs with all 7 sections
-    tab1, tab2, tab3, tab4, tab5, tab6, tab7 = st.tabs([
+    tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs([
         "Model Structure", "Comparison", "Model Attention", 
-        "Tokenization", "Embeddings", "Hardware", "Stats & Memory"
+        "Tokenization", "Hardware", "Memory"
     ])
     
     with tab1:
@@ -303,10 +206,6 @@ def main():
     with tab3:
         st.subheader("Model-specific Visualizations")
         visualize_attention_patterns()
-        if selected_model == "BERT":
-            st.write("BERT-specific visualization example")
-        elif selected_model == "GPT-2":
-            st.write("GPT-2 attention mask visualization")
     
     with tab4:
         st.subheader("📝 Tokenization Visualization")
@@ -317,7 +216,6 @@ def main():
             st.markdown("**Tokenized Output**")
             tokens = tokenizer.tokenize(input_text)
             st.write(tokens)
-            
         with col2:
             st.markdown("**Token IDs**")
             encoded_ids = tokenizer.encode(input_text)
@@ -328,15 +226,7 @@ def main():
             "Token": tokens,
             "ID": encoded_ids[1:-1] if tokenizer.cls_token else encoded_ids
         })
-        st.dataframe(
-            token_data,
-            height=150,
-            use_container_width=True,
-            column_config={
-                "Token": "Token",
-                "ID": {"header": "ID", "help": "Numerical representation of the token"}
-            }
-        )
+        st.dataframe(token_data, height=150, use_container_width=True)
         
         st.markdown(f"""
         **Tokenizer Info:**
@@ -345,16 +235,40 @@ def main():
         - Padding token: `{tokenizer.pad_token}`
         - Max length: `{tokenizer.model_max_length}`
         """)
-    
+
     with tab5:
-        embedding_projector()
-    
+        st.subheader("🖥️ Hardware Recommendation")
+        params = model_info["params"]
+        recommendation = get_hardware_recommendation(params)
+        
+        st.markdown(f"**Recommended hardware for {selected_model}:**")
+        st.info(recommendation)
+        
+        st.markdown("""
+        **Recommendation Criteria:**
+        - <100M parameters: Suitable for CPU or entry-level GPUs
+        - 100-200M parameters: Requires mid-range GPUs
+        - >200M parameters: Needs high-end GPUs/TPUs
+        """)
+
     with tab6:
-        hardware_recommendations(model_info)
-    
-    with tab7:
-        model_zoo_statistics()
-        memory_usage_estimator(model_info)
+        st.subheader("💾 Memory Usage Estimation")
+        params = model_info["params"]
+        memory_mb = params * 4  # 1M params ≈ 4MB in FP32
+        memory_gb = memory_mb / 1024
+        
+        st.metric("Estimated Memory (FP32)", 
+                 f"{memory_mb:.1f} MB / {memory_gb:.2f} GB")
+        
+        st.markdown("""
+        **Memory Notes:**
+        - Based on 4 bytes per parameter (FP32 precision)
+        - Actual usage varies with:
+          - Batch size
+          - Sequence length
+          - Precision (FP16/FP32)
+          - Optimizer states (training)
+        """)
 
 if __name__ == "__main__":
     main()