Upload viz_magnitude.py with huggingface_hub

viz_magnitude.py

@@ -0,0 +1,165 @@
+"""
+NeuroScope — Activation Magnitude Chart
+
+Bar chart showing L2 norm of hidden states across all layers.
+Highlights layers 9, 18, 27 (used by the Activation Avatars adapter).
+Shows how activation magnitude evolves through the network depth.
+
+All charts use Plotly with the project dark theme (#1a1a2e bg, #e6b800 accent).
+"""
+
+import numpy as np
+import plotly.graph_objects as go
+from extraction import ExtractionResult
+
+# ---------------------------------------------------------------------------
+# Theme constants
+# ---------------------------------------------------------------------------
+BG_COLOR = "#1a1a2e"
+PAPER_COLOR = "#1a1a2e"
+TEXT_COLOR = "#e0e0e0"
+ACCENT_COLOR = "#e6b800"
+GRID_COLOR = "#2a2a4e"
+BAR_COLOR = "#4a90d9"        # Default bar color
+HIGHLIGHT_COLOR = "#e6b800"  # Gold for highlighted layers
+EMBED_COLOR = "#7b68ee"      # Purple for embedding layer
+
+# Layers used by the Activation Avatars adapter
+AVATAR_LAYERS = {9, 18, 27}
+
+
+def create_magnitude_chart(
+    result: ExtractionResult,
+    metric: str = "mean_l2",
+) -> go.Figure:
+    """Create a bar chart of activation magnitude across all layers.
+
+    Args:
+        result: Extraction output containing hidden states.
+        metric: Magnitude metric to compute:
+            - "mean_l2": Mean L2 norm across tokens (default)
+            - "max_l2":  Max L2 norm across tokens
+            - "mean_abs": Mean absolute activation value
+
+    Returns:
+        Plotly Figure with interactive bar chart.
+    """
+    # hidden_states shape: (num_layers+1, seq_len, hidden_dim)
+    # Index 0 = embedding layer, 1..num_layers = transformer layers
+    hs = result.hidden_states
+    n_total = hs.shape[0]  # num_layers + 1
+
+    # Compute magnitude per layer
+    if metric == "mean_l2":
+        # L2 norm per token, then average across tokens
+        magnitudes = np.array([
+            np.linalg.norm(hs[i], axis=-1).mean() for i in range(n_total)
+        ])
+        y_label = "Mean L2 Norm"
+    elif metric == "max_l2":
+        magnitudes = np.array([
+            np.linalg.norm(hs[i], axis=-1).max() for i in range(n_total)
+        ])
+        y_label = "Max L2 Norm"
+    elif metric == "mean_abs":
+        magnitudes = np.array([
+            np.abs(hs[i]).mean() for i in range(n_total)
+        ])
+        y_label = "Mean |Activation|"
+    else:
+        raise ValueError(f"Unknown metric: {metric}")
+
+    # Build labels and colors
+    labels = ["Embed"] + [f"L{i}" for i in range(result.num_layers)]
+    colors = []
+    for i in range(n_total):
+        if i == 0:
+            colors.append(EMBED_COLOR)
+        elif (i - 1) in AVATAR_LAYERS:
+            colors.append(HIGHLIGHT_COLOR)
+        else:
+            colors.append(BAR_COLOR)
+
+    # Hover text with detailed info
+    hover_text = []
+    for i in range(n_total):
+        layer_name = "Embedding" if i == 0 else f"Layer {i - 1}"
+        tag = " ⭐ Avatar Layer" if (i > 0 and (i - 1) in AVATAR_LAYERS) else ""
+        hover_text.append(
+            f"{layer_name}{tag}<br>"
+            f"{y_label}: {magnitudes[i]:.2f}<br>"
+            f"Δ from prev: {magnitudes[i] - magnitudes[max(0, i-1)]:+.2f}"
+        )
+
+    fig = go.Figure(
+        data=go.Bar(
+            x=labels,
+            y=magnitudes,
+            marker_color=colors,
+            text=[f"{m:.1f}" for m in magnitudes],
+            textposition="outside",
+            textfont=dict(color=TEXT_COLOR, size=8),
+            hovertext=hover_text,
+            hoverinfo="text",
+        )
+    )
+
+    # Add subtle markers for avatar layers
+    avatar_x = []
+    avatar_y = []
+    for layer in sorted(AVATAR_LAYERS):
+        idx = layer + 1  # +1 for embedding offset
+        if idx < n_total:
+            avatar_x.append(labels[idx])
+            avatar_y.append(magnitudes[idx])
+
+    if avatar_x:
+        fig.add_trace(
+            go.Scatter(
+                x=avatar_x,
+                y=avatar_y,
+                mode="markers",
+                marker=dict(
+                    symbol="star",
+                    size=12,
+                    color=HIGHLIGHT_COLOR,
+                    line=dict(width=1, color="white"),
+                ),
+                name="Avatar Layers",
+                hoverinfo="skip",
+                showlegend=True,
+            )
+        )
+
+    fig.update_layout(
+        title=dict(
+            text=f"Activation Magnitude by Layer ({y_label})",
+            font=dict(color=ACCENT_COLOR, size=14),
+        ),
+        xaxis=dict(
+            title=dict(text="Layer", font=dict(color=TEXT_COLOR, size=11)),
+            tickfont=dict(color=TEXT_COLOR, size=8),
+            gridcolor=GRID_COLOR,
+            tickangle=45,
+        ),
+        yaxis=dict(
+            title=dict(text=y_label, font=dict(color=TEXT_COLOR, size=11)),
+            tickfont=dict(color=TEXT_COLOR, size=9),
+            gridcolor=GRID_COLOR,
+            zeroline=False,
+        ),
+        paper_bgcolor=PAPER_COLOR,
+        plot_bgcolor=BG_COLOR,
+        margin=dict(l=60, r=20, t=50, b=60),
+        height=480,
+        showlegend=True,
+        legend=dict(
+            font=dict(color=TEXT_COLOR, size=10),
+            bgcolor="rgba(26, 26, 46, 0.8)",
+            x=0.02,
+            y=0.98,
+        ),
+        bargap=0.15,
+    )
+
+    return fig