added full screen

app.py

@@ -1,8 +1,3 @@
-# FULL LLM VISUALIZER — OPTION A (ADVANCED)
-# stable + patched + safe for HuggingFace Spaces (CPU or GPU)
-# recommended models: distilgpt2, gpt2
-# author: ChatGPT
-
 import gradio as gr
 import torch
 import numpy as np
@@ -12,11 +7,74 @@ import pandas as pd
 from sklearn.decomposition import PCA
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import html
+import time
 
 DEFAULT_MODEL = "distilgpt2"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL_CACHE = {}
 
+# ---------------- Fullscreen helper (injected JS per plot) ----------------
+
+def fullscreen_plot_js(plot_id):
+    # returns JS + function named openFull_{plot_id}()
+    # clicking overlay removes it
+    safe_id = plot_id.replace("-", "_")
+    return f"""
+    <script>
+    function openFull_{safe_id}() {{
+        // find the plot element (img or canvas) inside the element with elem_id
+        const container = document.getElementById("{plot_id}");
+        if (!container) {{
+            return;
+        }}
+        const img = container.querySelector("img, canvas");
+        if (!img) {{
+            // no canvas or img yet, inform user
+            const hint = document.createElement('div');
+            hint.style.position='fixed';
+            hint.style.bottom='20px';
+            hint.style.left='20px';
+            hint.style.padding='8px 12px';
+            hint.style.background='#222';
+            hint.style.color='white';
+            hint.style.borderRadius='8px';
+            hint.style.zIndex = 99999;
+            hint.innerText = 'Plot not rendered yet. Please run analysis first.';
+            document.body.appendChild(hint);
+            setTimeout(()=>hint.remove(), 2500);
+            return;
+        }}
+        const modal = document.createElement('div');
+        modal.style.position = 'fixed';
+        modal.style.top = '0';
+        modal.style.left = '0';
+        modal.style.width = '100vw';
+        modal.style.height = '100vh';
+        modal.style.background = 'rgba(0,0,0,0.88)';
+        modal.style.zIndex = '999999';
+        modal.style.display = 'flex';
+        modal.style.alignItems = 'center';
+        modal.style.justifyContent = 'center';
+        modal.onclick = () => modal.remove();
+
+        const clone = img.cloneNode(true);
+        clone.style.maxWidth = '95%';
+        clone.style.maxHeight = '95%';
+        clone.style.border = '2px solid rgba(255,255,255,0.85)';
+        clone.style.borderRadius = '10px';
+        clone.style.boxShadow = '0 8px 40px rgba(0,0,0,0.8)';
+        modal.appendChild(clone);
+
+        document.body.appendChild(modal);
+    }}
+    </script>
+    """
+
+def fullscreen_button_html(plot_id, label="🔍 Full Screen"):
+    # wrapper HTML: JS function + button
+    return fullscreen_plot_js(plot_id) + f'<button onclick="openFull_{plot_id.replace("-","_")}()" style="margin-top:6px;padding:6px 10px;border-radius:8px;border:1px solid #ddd;background:white;">{label}</button>'
+
+
 # ---------------- CORE UTILS ----------------
 
 def load_model(model_name):
@@ -197,8 +255,9 @@ def compute_residuals_safe(model, inputs):
                 mlp_out = block.mlp(ln2)
                 x = x + mlp_out
 
-                attn_norms.append(float(torch.norm(attn_out).cpu()))
-                mlp_norms.append(float(torch.norm(mlp_out).cpu()))
+                # detach to avoid requires_grad warning
+                attn_norms.append(float(torch.norm(attn_out.detach()).cpu()))
+                mlp_norms.append(float(torch.norm(mlp_out.detach()).cpu()))
             except:
                 # fallback safe zero
                 attn_norms.append(0.0)
@@ -311,15 +370,19 @@ with gr.Blocks(title="LLM Visualizer — Full", theme=gr.themes.Soft()) as demo:
             layer_slider = gr.Slider(0, 0, value=0, step=1, label="Layer")
             head_slider = gr.Slider(0, 0, value=0, step=1, label="Head")
             token_step = gr.Slider(0, 0, value=0, step=1, label="Token index")
-            attn_plot = gr.Plot()
-
+            attn_plot = gr.Plot(elem_id="attn_plot")
+            attn_fs = gr.HTML(fullscreen_button_html("attn_plot"))
         with gr.Column():
-            pca_plot = gr.Plot()
-            step_attn_plot = gr.Plot()
-            probs_plot = gr.Plot()
+            pca_plot = gr.Plot(elem_id="pca_plot")
+            pca_fs = gr.HTML(fullscreen_button_html("pca_plot"))
+            step_attn_plot = gr.Plot(elem_id="step_attn_plot")
+            step_fs = gr.HTML(fullscreen_button_html("step_attn_plot"))
+            probs_plot = gr.Plot(elem_id="probs_plot")
+            probs_fs = gr.HTML(fullscreen_button_html("probs_plot"))
 
     # Panel 3 — Residuals
-    residual_plot = gr.Plot()
+    residual_plot = gr.Plot(elem_id="residual_plot")
+    residual_fs = gr.HTML(fullscreen_button_html("residual_plot"))
 
     # Panel 4 — Neuron explorer
     with gr.Row():
@@ -334,7 +397,8 @@ with gr.Blocks(title="LLM Visualizer — Full", theme=gr.themes.Soft()) as demo:
         patch_from = gr.Slider(0, 0, value=0, step=1, label="Copy from position")
         patch_scale = gr.Number(label="Scale", value=1.0)
         patch_btn = gr.Button("Run patch")
-        patch_output = gr.Plot()
+        patch_output = gr.Plot(elem_id="patch_plot")
+        patch_fs = gr.HTML(fullscreen_button_html("patch_plot"))
 
     state = gr.State()
 
@@ -358,7 +422,9 @@ with gr.Blocks(title="LLM Visualizer — Full", theme=gr.themes.Soft()) as demo:
                 patch_layer: gr.update(maximum=0),
                 patch_pos: gr.update(maximum=0),
                 patch_from: gr.update(maximum=0),
-                state: res
+                state: res,
+                step_attn_plot: gr.update(value=None),
+                patch_output: gr.update(value=None),
             }
 
         tokens = res["tokens"]
@@ -397,7 +463,9 @@ with gr.Blocks(title="LLM Visualizer — Full", theme=gr.themes.Soft()) as demo:
             patch_layer: gr.update(maximum=L-1, value=0),
             patch_pos: gr.update(maximum=T, value=0),
             patch_from: gr.update(maximum=T, value=0),
-            state: res
+            state: res,
+            step_attn_plot: gr.update(value=None),
+            patch_output: gr.update(value=None),
         }
 
 
@@ -410,7 +478,7 @@ with gr.Blocks(title="LLM Visualizer — Full", theme=gr.themes.Soft()) as demo:
             layer_slider, head_slider, token_step,
             residual_plot, neuron_table,
             patch_layer, patch_pos, patch_from,
-            state
+            state, step_attn_plot, patch_output
         ]
     )
 
@@ -494,4 +562,4 @@ with gr.Blocks(title="LLM Visualizer — Full", theme=gr.themes.Soft()) as demo:
                     [state, patch_layer, patch_pos, patch_from, patch_scale, model_name],
                     [patch_output])
 
-demo.launch()
+demo.launch()