Enhance Neuroscope UI with improved styling and layout

app.py

@@ -23,7 +23,25 @@ def calculate_color(val, max_val, min_val):
 
 style_string = """<style>
     span.token {
-        border: 1px solid rgb(123, 123, 123)
+        border: 1px solid #ddd;
+        padding: 2px 4px;
+        margin: 1px;
+        border-radius: 3px;
+        font-family: monospace;
+    }
+    .neuron-info {
+        background: #f5f5f5;
+        padding: 15px;
+        border-radius: 8px;
+        margin-bottom: 15px;
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    }
+    .token-container {
+        line-height: 2;
+        background: white;
+        padding: 20px;
+        border-radius: 8px;
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
     }
 </style>"""
 
@@ -41,19 +59,24 @@ def basic_neuron_vis(text, layer, neuron_index, max_val=None, min_val=None):
         min_val = act_min
     
     htmls = [style_string]
-    htmls.append(f"<h4>Layer: <b>{layer}</b>. Neuron Index: <b>{neuron_index}</b></h4>")
-    htmls.append(f"<h4>Max Range: <b>{max_val:.4f}</b>. Min Range: <b>{min_val:.4f}</b></h4>")
+    htmls.append('<div class="neuron-info">')
+    htmls.append(f"<h3 style='margin-top: 0;'>Neuron Information</h3>")
+    htmls.append(f"<p>Layer: <b>{layer}</b> | Neuron Index: <b>{neuron_index}</b></p>")
+    htmls.append(f"<p>Max Range: <b>{max_val:.4f}</b> | Min Range: <b>{min_val:.4f}</b></p>")
     
     if act_max != max_val or act_min != min_val:
         htmls.append(
-            f"<h4>Custom Range Set. Max Act: <b>{act_max:.4f}</b>. Min Act: <b>{act_min:.4f}</b></h4>"
+            f"<p>Custom Range Set | Max Act: <b>{act_max:.4f}</b> | Min Act: <b>{act_min:.4f}</b></p>"
         )
+    htmls.append('</div>')
     
+    htmls.append('<div class="token-container">')
     str_tokens = model.to_str_tokens(text)
     for tok, act in zip(str_tokens, acts):
         htmls.append(
             f"<span class='token' style='background-color:{calculate_color(act, max_val, min_val)}' >{tok}</span>"
         )
+    htmls.append('</div>')
     
     return "".join(htmls)
 
@@ -123,21 +146,48 @@ default_neuron_index = 1
 default_max_val = 4.0
 default_min_val = 0.0
 
-with gr.Blocks() as demo:
-    gr.HTML(value=f"Neuroscope for {model_name}")
+with gr.Blocks(css="footer {display: none}") as demo:
+    gr.HTML(
+        value=f"<h1 style='text-align: center; margin-bottom: 1rem;'>Neuroscope for {model_name}</h1>"
+    )
     with gr.Row():
-        with gr.Column():
-            text = gr.Textbox(label="Text", value=default_text)
-            layer = gr.Number(label="Layer", value=default_layer, precision=0)
-            neuron_index = gr.Number(
-                label="Neuron Index", value=default_neuron_index, precision=0
-            )
-            max_val = gr.Number(label="Max Value", value=default_max_val)
-            min_val = gr.Number(label="Min Value", value=default_min_val)
-            inputs = [text, layer, neuron_index, max_val, min_val]
-        with gr.Column():
+        with gr.Column(scale=1):
+            with gr.Box():
+                gr.HTML("<h3>Controls</h3>")
+                text = gr.Textbox(
+                    label="Input Text",
+                    value=default_text,
+                    lines=10
+                )
+                with gr.Row():
+                    layer = gr.Number(
+                        label="Layer",
+                        value=default_layer,
+                        precision=0,
+                        container=True
+                    )
+                    neuron_index = gr.Number(
+                        label="Neuron Index",
+                        value=default_neuron_index,
+                        precision=0,
+                        container=True
+                    )
+                with gr.Row():
+                    max_val = gr.Number(
+                        label="Max Value",
+                        value=default_max_val,
+                        container=True
+                    )
+                    min_val = gr.Number(
+                        label="Min Value",
+                        value=default_min_val,
+                        container=True
+                    )
+                inputs = [text, layer, neuron_index, max_val, min_val]
+        
+        with gr.Column(scale=1):
             out = gr.HTML(
-                label="Neuron Acts",
+                label="Neuron Activation Visualization",
                 value=basic_neuron_vis(
                     default_text,
                     default_layer,
@@ -146,6 +196,7 @@ with gr.Blocks() as demo:
                     default_min_val,
                 ),
             )
+
     for inp in inputs:
         inp.change(basic_neuron_vis, inputs, out)