Add PLV matrix heatmap to Inspect tab

app.py

@@ -14,6 +14,7 @@ from pathlib import Path
 from sklearn.decomposition import PCA
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import plotly.graph_objects as go
+import plotly.express as px
 import spaces  # For ZeroGPU on Hugging Face
 
 # --- CONFIG ---
@@ -255,7 +256,32 @@ def analyze_text(text):
     else:
         interpretation = "**Balanced**  \nMixed patterns - both structure and meaning equally present"
 
-    return fig, interpretation, score
+    # Create PLV matrix heatmap (reshape to 256x256)
+    plv_np = plv_pred[0].cpu().numpy()
+    plv_matrix = plv_np[:65536].reshape(256, 256)
+
+    fig_plv = px.imshow(
+        plv_matrix,
+        color_continuous_scale='Viridis',
+        aspect='auto'
+    )
+    fig_plv.update_layout(
+        coloraxis_showscale=True,
+        coloraxis=dict(
+            colorbar=dict(
+                thickness=10,
+                len=0.7,
+                title=dict(text="Synchrony", side="right"),
+                tickfont=dict(size=10)
+            )
+        ),
+        margin={'l': 0, 'r': 40, 't': 10, 'b': 0},
+        height=300
+    )
+    fig_plv.update_xaxes(visible=False)
+    fig_plv.update_yaxes(visible=False)
+
+    return fig, interpretation, score, fig_plv
 
 @spaces.GPU(duration=60)
 def generate_steered(prompt, alpha, max_tokens):
@@ -523,14 +549,22 @@ with gr.Blocks(title="Cognitive Proxy") as demo:
                         - **Range** typically -300 to +300
                         """)
 
+                    with gr.Accordion("View brain connectivity pattern", open=False):
+                        gr.Markdown("""
+                        Phase-Locking Value (PLV) shows how synchronized different brain regions are.
+                        Brighter colors = stronger synchronization between sensor pairs.
+                        Each pixel represents connectivity between two of 256 MEG sensors.
+                        """)
+                        plv_plot = gr.Plot(label="")
+
             def analyze_text_wrapper(text):
-                fig, interp, _ = analyze_text(text)  # Ignore the score
-                return fig, interp
+                fig, interp, _, fig_plv = analyze_text(text)
+                return fig, interp, fig_plv
 
             analyze_btn.click(
                 analyze_text_wrapper,
                 inputs=[text_input],
-                outputs=[gauge_plot, interpretation]
+                outputs=[gauge_plot, interpretation, plv_plot]
             )
 
         # Steer Tab