Upload viz_comparison.py with huggingface_hub

viz_comparison.py

@@ -0,0 +1,366 @@
+"""
+NeuroScope — Comparison Mode Visualizations (Phase 2)
+
+Side-by-side and overlay visualizations comparing activation patterns
+between two different prompts. Uses joint normalization so the two
+prompts are visually comparable.
+
+All charts use Plotly with the project dark theme (#1a1a2e bg, #e6b800 accent).
+"""
+
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+from extraction import ExtractionResult
+
+# ---------------------------------------------------------------------------
+# Theme constants
+# ---------------------------------------------------------------------------
+BG_COLOR = "#1a1a2e"
+PAPER_COLOR = "#1a1a2e"
+TEXT_COLOR = "#e0e0e0"
+ACCENT_COLOR = "#e6b800"
+GRID_COLOR = "#2a2a4e"
+
+COLOR_A = "#e6b800"  # Gold for Prompt A
+COLOR_B = "#4a90d9"  # Blue for Prompt B
+
+ATTN_COLORSCALE = [
+    [0.0, "#0d0d1a"],
+    [0.15, "#1a1a3e"],
+    [0.3, "#2e1a00"],
+    [0.5, "#7a5500"],
+    [0.7, "#b38600"],
+    [0.85, "#e6b800"],
+    [1.0, "#ffd633"],
+]
+
+TOKEN_LAYER_COLORSCALE = [
+    [0.0, "#0d0d1a"],
+    [0.1, "#1a1040"],
+    [0.25, "#2d1b69"],
+    [0.4, "#5e2d8e"],
+    [0.55, "#8e4585"],
+    [0.7, "#c46a3a"],
+    [0.85, "#e6b800"],
+    [1.0, "#ffd633"],
+]
+
+# Diverging colorscale for difference maps (blue-white-red)
+DIFF_COLORSCALE = [
+    [0.0, "#2166ac"],
+    [0.25, "#67a9cf"],
+    [0.5, "#1a1a2e"],
+    [0.75, "#ef8a62"],
+    [1.0, "#b2182b"],
+]
+
+
+# ---------------------------------------------------------------------------
+# Attention comparison
+# ---------------------------------------------------------------------------
+def create_attention_comparison(
+    result_a: ExtractionResult,
+    result_b: ExtractionResult,
+    layer: int = 0,
+    head: str = "average",
+) -> go.Figure:
+    """Side-by-side attention heatmaps with difference overlay."""
+    layer = int(np.clip(layer, 0, min(result_a.num_layers, result_b.num_layers) - 1))
+
+    def _get_matrix(result, layer, head):
+        attn = result.attentions[layer]
+        if head == "average":
+            return attn.mean(axis=0)
+        elif head == "max":
+            return attn.max(axis=0)
+        else:
+            h = int(np.clip(int(head), 0, result.num_heads - 1))
+            return attn[h]
+
+    mat_a = _get_matrix(result_a, layer, head)
+    mat_b = _get_matrix(result_b, layer, head)
+    labels_a = [t[:10] for t in result_a.tokens]
+    labels_b = [t[:10] for t in result_b.tokens]
+
+    head_label = f"avg" if head == "average" else f"max" if head == "max" else f"h{head}"
+
+    fig = make_subplots(
+        rows=1, cols=2,
+        subplot_titles=[f"Prompt A — L{layer} {head_label}", f"Prompt B — L{layer} {head_label}"],
+        horizontal_spacing=0.08,
+    )
+
+    # Joint scale
+    vmax = max(float(mat_a.max()), float(mat_b.max()))
+    if vmax == 0:
+        vmax = 1.0
+
+    def _hover(tokens, mat):
+        sl = len(tokens)
+        h = np.empty((sl, sl), dtype=object)
+        for i in range(sl):
+            for j in range(sl):
+                h[i, j] = f"From: {tokens[i]}<br>To: {tokens[j]}<br>Wt: {mat[i, j]:.4f}"
+        return h
+
+    fig.add_trace(
+        go.Heatmap(
+            z=mat_a, x=labels_a, y=labels_a,
+            text=_hover(result_a.tokens, mat_a), hoverinfo="text",
+            colorscale=ATTN_COLORSCALE, zmin=0, zmax=vmax, showscale=False,
+        ), row=1, col=1,
+    )
+    fig.add_trace(
+        go.Heatmap(
+            z=mat_b, x=labels_b, y=labels_b,
+            text=_hover(result_b.tokens, mat_b), hoverinfo="text",
+            colorscale=ATTN_COLORSCALE, zmin=0, zmax=vmax,
+            colorbar=dict(
+                title=dict(text="Attn", font=dict(color=TEXT_COLOR)),
+                tickfont=dict(color=TEXT_COLOR),
+            ),
+        ), row=1, col=2,
+    )
+
+    fig.update_layout(
+        paper_bgcolor=PAPER_COLOR, plot_bgcolor=BG_COLOR,
+        height=480, margin=dict(l=60, r=30, t=50, b=60),
+    )
+    for col in (1, 2):
+        fig.update_xaxes(tickfont=dict(color=TEXT_COLOR, size=8), tickangle=45, row=1, col=col)
+        fig.update_yaxes(
+            tickfont=dict(color=TEXT_COLOR, size=8), autorange="reversed",
+            row=1, col=col,
+        )
+    for ann in fig.layout.annotations:
+        ann.font = dict(color=ACCENT_COLOR, size=12)
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Magnitude comparison
+# ---------------------------------------------------------------------------
+def create_magnitude_comparison(
+    result_a: ExtractionResult,
+    result_b: ExtractionResult,
+    metric: str = "mean_l2",
+) -> go.Figure:
+    """Overlaid magnitude bar chart for two prompts."""
+
+    def _compute_mag(hs, metric):
+        n = hs.shape[0]
+        if metric == "mean_l2":
+            return np.array([np.linalg.norm(hs[i], axis=-1).mean() for i in range(n)])
+        elif metric == "max_l2":
+            return np.array([np.linalg.norm(hs[i], axis=-1).max() for i in range(n)])
+        else:  # mean_abs
+            return np.array([np.abs(hs[i]).mean() for i in range(n)])
+
+    mag_a = _compute_mag(result_a.hidden_states, metric)
+    mag_b = _compute_mag(result_b.hidden_states, metric)
+    labels = ["Emb"] + [f"L{i}" for i in range(result_a.num_layers)]
+
+    fig = go.Figure()
+
+    fig.add_trace(go.Bar(
+        x=labels, y=mag_a,
+        name="Prompt A", marker_color=COLOR_A, opacity=0.8,
+        hovertext=[f"Prompt A<br>{labels[i]}<br>{metric}: {mag_a[i]:.2f}" for i in range(len(labels))],
+        hoverinfo="text",
+    ))
+    fig.add_trace(go.Bar(
+        x=labels, y=mag_b,
+        name="Prompt B", marker_color=COLOR_B, opacity=0.8,
+        hovertext=[f"Prompt B<br>{labels[i]}<br>{metric}: {mag_b[i]:.2f}" for i in range(len(labels))],
+        hoverinfo="text",
+    ))
+
+    # Difference line
+    diff = mag_a[:len(mag_b)] - mag_b[:len(mag_a)]
+    min_len = min(len(mag_a), len(mag_b))
+    fig.add_trace(go.Scatter(
+        x=labels[:min_len], y=diff[:min_len],
+        name="Δ (A − B)", mode="lines+markers",
+        line=dict(color="#e05050", width=2, dash="dot"),
+        marker=dict(size=4, color="#e05050"),
+        yaxis="y2",
+        hovertext=[f"Δ at {labels[i]}: {diff[i]:+.2f}" for i in range(min_len)],
+        hoverinfo="text",
+    ))
+
+    fig.update_layout(
+        title=dict(text=f"Activation Magnitude Comparison ({metric})", font=dict(color=ACCENT_COLOR, size=14)),
+        barmode="group",
+        paper_bgcolor=PAPER_COLOR, plot_bgcolor=BG_COLOR,
+        height=480, margin=dict(l=60, r=60, t=50, b=60),
+        xaxis=dict(tickfont=dict(color=TEXT_COLOR, size=8), gridcolor=GRID_COLOR, tickangle=45),
+        yaxis=dict(
+            title=dict(text=metric, font=dict(color=TEXT_COLOR, size=11)),
+            tickfont=dict(color=TEXT_COLOR, size=9), gridcolor=GRID_COLOR,
+        ),
+        yaxis2=dict(
+            title=dict(text="Δ", font=dict(color="#e05050", size=11)),
+            tickfont=dict(color="#e05050", size=9),
+            overlaying="y", side="right", zeroline=True,
+            zerolinecolor="rgba(224,80,80,0.3)",
+        ),
+        legend=dict(font=dict(color=TEXT_COLOR, size=10), bgcolor="rgba(26,26,46,0.8)"),
+        bargap=0.15,
+    )
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Token-Layer grid comparison
+# ---------------------------------------------------------------------------
+def create_token_layer_comparison(
+    result_a: ExtractionResult,
+    result_b: ExtractionResult,
+    normalize: str = "global",
+) -> go.Figure:
+    """Side-by-side token-layer activation grids."""
+    norms_a = np.linalg.norm(result_a.hidden_states, axis=-1)
+    norms_b = np.linalg.norm(result_b.hidden_states, axis=-1)
+
+    # Joint normalization
+    if normalize == "global":
+        vmin = min(float(norms_a.min()), float(norms_b.min()))
+        vmax = max(float(norms_a.max()), float(norms_b.max()))
+        if vmax > vmin:
+            disp_a = (norms_a - vmin) / (vmax - vmin)
+            disp_b = (norms_b - vmin) / (vmax - vmin)
+        else:
+            disp_a, disp_b = norms_a * 0, norms_b * 0
+    else:
+        disp_a, disp_b = norms_a.copy(), norms_b.copy()
+
+    labels_a = [t[:10] for t in result_a.tokens]
+    labels_b = [t[:10] for t in result_b.tokens]
+    y_labels = ["Emb"] + [f"L{i}" for i in range(result_a.num_layers)]
+
+    fig = make_subplots(
+        rows=1, cols=2,
+        subplot_titles=["Prompt A — Token×Layer", "Prompt B — Token×Layer"],
+        horizontal_spacing=0.08,
+    )
+
+    def _hover(tokens, norms):
+        h = np.empty(norms.shape, dtype=object)
+        for i in range(norms.shape[0]):
+            lname = "Embedding" if i == 0 else f"Layer {i-1}"
+            for j in range(len(tokens)):
+                h[i, j] = f"{tokens[j]}<br>{lname}<br>L2: {norms[i, j]:.2f}"
+        return h
+
+    fig.add_trace(
+        go.Heatmap(
+            z=disp_a, x=labels_a, y=y_labels,
+            text=_hover(result_a.tokens, norms_a), hoverinfo="text",
+            colorscale=TOKEN_LAYER_COLORSCALE, showscale=False,
+        ), row=1, col=1,
+    )
+    fig.add_trace(
+        go.Heatmap(
+            z=disp_b, x=labels_b, y=y_labels,
+            text=_hover(result_b.tokens, norms_b), hoverinfo="text",
+            colorscale=TOKEN_LAYER_COLORSCALE, showscale=True,
+            colorbar=dict(
+                title=dict(text="Norm", font=dict(color=TEXT_COLOR)),
+                tickfont=dict(color=TEXT_COLOR),
+            ),
+        ), row=1, col=2,
+    )
+
+    fig.update_layout(
+        paper_bgcolor=PAPER_COLOR, plot_bgcolor=BG_COLOR,
+        height=520, margin=dict(l=60, r=30, t=50, b=30),
+    )
+    for col in (1, 2):
+        fig.update_xaxes(tickfont=dict(color=TEXT_COLOR, size=8), tickangle=45, side="top", row=1, col=col)
+        fig.update_yaxes(tickfont=dict(color=TEXT_COLOR, size=7), autorange="reversed", row=1, col=col)
+    for ann in fig.layout.annotations:
+        ann.font = dict(color=ACCENT_COLOR, size=12)
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Scatter comparison (overlay both prompts)
+# ---------------------------------------------------------------------------
+def create_scatter_comparison(
+    result_a: ExtractionResult,
+    result_b: ExtractionResult,
+    layer: int = 18,
+    method: str = "pca",
+) -> go.Figure:
+    """Overlay scatter plot with both prompts' tokens in same reduced space."""
+    from viz_scatter import _run_pca, _run_umap
+
+    reduce_fn = _run_umap if method == "umap" else _run_pca
+    hs_idx = int(np.clip(layer + 1, 0, result_a.hidden_states.shape[0] - 1))
+
+    data_a = result_a.hidden_states[hs_idx].astype(np.float64)
+    data_b = result_b.hidden_states[
+        int(np.clip(layer + 1, 0, result_b.hidden_states.shape[0] - 1))
+    ].astype(np.float64)
+
+    # Joint reduction for fair comparison
+    stacked = np.vstack([data_a, data_b])
+    reduced = reduce_fn(stacked)
+    n_a = len(result_a.tokens)
+    red_a = reduced[:n_a]
+    red_b = reduced[n_a:]
+
+    fig = go.Figure()
+
+    # Prompt A tokens
+    fig.add_trace(go.Scatter(
+        x=red_a[:, 0], y=red_a[:, 1],
+        mode="markers+text",
+        marker=dict(size=12, color=COLOR_A, opacity=0.9, line=dict(width=1, color="white")),
+        text=[t[:8] for t in result_a.tokens],
+        textposition="top center",
+        textfont=dict(color=COLOR_A, size=9),
+        name="Prompt A",
+        hovertext=[
+            f"A: {result_a.tokens[j]}<br>Pos: {j}<br>x: {red_a[j,0]:.3f}, y: {red_a[j,1]:.3f}"
+            for j in range(n_a)
+        ],
+        hoverinfo="text",
+    ))
+
+    # Prompt B tokens
+    fig.add_trace(go.Scatter(
+        x=red_b[:, 0], y=red_b[:, 1],
+        mode="markers+text",
+        marker=dict(size=12, color=COLOR_B, opacity=0.9, symbol="diamond",
+                    line=dict(width=1, color="white")),
+        text=[t[:8] for t in result_b.tokens],
+        textposition="bottom center",
+        textfont=dict(color=COLOR_B, size=9),
+        name="Prompt B",
+        hovertext=[
+            f"B: {result_b.tokens[j]}<br>Pos: {j}<br>x: {red_b[j,0]:.3f}, y: {red_b[j,1]:.3f}"
+            for j in range(len(result_b.tokens))
+        ],
+        hoverinfo="text",
+    ))
+
+    fig.update_layout(
+        title=dict(
+            text=f"Token Space Comparison ({method.upper()}) — Layer {layer}",
+            font=dict(color=ACCENT_COLOR, size=14),
+        ),
+        xaxis=dict(
+            title=dict(text=f"{method.upper()} 1", font=dict(color=TEXT_COLOR, size=11)),
+            tickfont=dict(color=TEXT_COLOR, size=9), gridcolor=GRID_COLOR, zeroline=False,
+        ),
+        yaxis=dict(
+            title=dict(text=f"{method.upper()} 2", 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,
+        height=480, margin=dict(l=60, r=30, t=50, b=50),
+        legend=dict(font=dict(color=TEXT_COLOR, size=10), bgcolor="rgba(26,26,46,0.8)"),
+    )
+    return fig