Update app.py

app.py

@@ -34,6 +34,29 @@ client = Client(f"http://{os.environ['SERVER']}/predict")
 def get_layerwise_nonlinearity(task_name: str, model_name: str, text: str, normalization_type: str) -> Tuple[Any, str]:
     return client.send_request(task_name, model_name, text, normalization_type)
 
+def update_output(task_name: str, model_name: str, text: str, normalization_type: str) -> Tuple[Any]:
+    img, _ = get_layerwise_nonlinearity(task_name, model_name, text, normalization_type)
+    return img
+
+def set_default(task_name: str) -> str:
+    if task_name in ["Layer wise non-linearity", "Next-token prediction from intermediate representations", "Tokenwise loss without i-th layer"]:
+        return "token-wise"
+    return "global"
+
+def check_normalization(task_name: str, normalization_name) -> Tuple[str]:
+    if task_name == "Contextualization measurement" and normalization_name == "token-wise":
+        return "global"
+    return normalization_name
+
+def update_description(task_name: str) -> str:
+    descriptions = {
+        "Layer wise non-linearity": "Non-linearity per layer: shows how complex each layer's transformation is. Red = more nonlinear.",
+        "Next-token prediction from intermediate representations": "Layerwise token prediction: when does the model start guessing correctly?",
+        "Contextualization measurement": "Context stored in each token: how well can the model reconstruct the previous context?",
+        "Layerwise predictions (logit lens)": "Logit lens: what does each layer believe the next token should be?",
+        "Tokenwise loss without i-th layer": "Layer ablation: how much does performance drop if a layer is removed?"
+    }
+    return descriptions.get(task_name, "ℹ️ No description available.")
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🔬 LLM-Microscope — Understanding Token Representations in Transformers")
@@ -65,22 +88,24 @@ with gr.Blocks() as demo:
             label="Select Normalization"
         )
 
+    task_description = gr.Markdown("ℹ️ Choose a mode to see what it does.")
+    
     with gr.Column():
         text_message = gr.Textbox(label="Enter your input text:", value="I love to live my life")
         submit = gr.Button("Submit")
         box_for_plot = gr.Image(label="Visualization", type="pil")
 
-        # 💬 Explanation below the visualization
-        explanation_text = gr.Markdown("""
-### 📘 Legend and Interpretation
-
+        with gr.Accordion("📘 Full Legend and Interpretation", open=False):
+            gr.Markdown("""
 This heatmap shows **how each token is processed** across layers of a language model. Here's how to read it:
 
 - **Rows**: layers of the model (bottom = deeper)
 - **Columns**: input tokens
 - **Colors**: intensity of effect (depends on the selected metric)
 
-**Metrics explained:**
+---
+
+### Metrics explained:
 
 - `Layer wise non-linearity`: how nonlinear the transformation is at each layer (red = more nonlinear).
 - `Next-token prediction from intermediate representations`: shows which layers begin to make good predictions.
@@ -89,29 +114,16 @@ This heatmap shows **how each token is processed** across layers of a language m
 - `Tokenwise loss without i-th layer`: shows how much each token depends on a specific layer. Red means performance drops if we skip this layer.
 
 Use this tool to **peek inside the black box** — it reveals which layers matter most, which tokens carry the most memory, and how LLMs evolve their predictions.
-        """)
-
-        def update_output(task_name: str, model_name: str, text: str, normalization_type: str) -> Tuple[Any]:
-            img, _ = get_layerwise_nonlinearity(task_name, model_name, text, normalization_type)
-            return img
-
-        def set_default(task_name: str) -> str:
-            if task_name in ["Layer wise non-linearity", "Next-token prediction from intermediate representations", "Tokenwise loss without i-th layer"]:
-                return "token-wise"
-            return "global"
-
-        def check_normalization(task_name: str, normalization_name) -> Tuple[str]:
-            if task_name == "Contextualization measurement" and normalization_name == "token-wise":
-                return "global"
-            return normalization_name
-
-        task_selector.select(set_default, [task_selector], [normalization_selector])
-        normalization_selector.select(check_normalization, [task_selector, normalization_selector], [normalization_selector])
-        submit.click(
-            fn=update_output, 
-            inputs=[task_selector, model_selector, text_message, normalization_selector], 
-            outputs=[box_for_plot]
-        )
+""")
+
+    task_selector.change(fn=update_description, inputs=[task_selector], outputs=[task_description])
+    task_selector.select(set_default, [task_selector], [normalization_selector])
+    normalization_selector.select(check_normalization, [task_selector, normalization_selector], [normalization_selector])
+    submit.click(
+        fn=update_output, 
+        inputs=[task_selector, model_selector, text_message, normalization_selector], 
+        outputs=[box_for_plot]
+    )
 
 if __name__ == "__main__":
     demo.launch(share=True, server_port=7860, server_name="0.0.0.0")