app.py

from __future__ import annotations

from pathlib import Path

import gradio as gr
import pandas as pd
import plotly.graph_objects as go

DATA = pd.read_parquet(
    Path(__file__).resolve().parent / "data" / "feature_exemplars.parquet"
)
FEATURES = sorted(DATA["feature"].unique().tolist())


def inspect_feature(feature: int) -> tuple[go.Figure, dict]:
    feature = int(feature)
    rows = DATA[DATA["feature"] == feature].sort_values("rank")
    figure = go.Figure(
        go.Bar(
            x=rows["token"].tolist(),
            y=rows["activation"].tolist(),
            marker_color="#a78bfa",
            customdata=rows["token_id"].tolist(),
            hovertemplate="token=%{x}<br>id=%{customdata}<br>activation=%{y:.3f}",
        )
    )
    figure.update_layout(
        title=f"Feature {feature}: top held-out token activations",
        xaxis_title="BPE token",
        yaxis_title="Sparse feature activation",
        template="plotly_dark",
    )
    return figure, {
        "feature": feature,
        "heldout_firing_rate": round(float(rows["firing_rate"].iloc[0]), 6),
        "top_tokens": rows["token"].tolist(),
        "warning": "Token exemplars are clues, not causal concept labels.",
    }


with gr.Blocks(title="SNIP Scope") as demo:
    gr.Markdown(
        "# SNIP Scope\n"
        "Explore top held-out token activations for 384 sparse features learned "
        "from the SNIP transformer's final hidden layer."
    )
    feature = gr.Slider(0, max(FEATURES), 0, step=1, label="Sparse feature")
    run = gr.Button("Inspect feature", variant="primary")
    exemplars = gr.Plot()
    details = gr.JSON()
    run.click(inspect_feature, feature, [exemplars, details])
    demo.load(inspect_feature, feature, [exemplars, details])


if __name__ == "__main__":
    demo.launch()