app.py

import gradio as gr
import torch
import numpy as np
import pandas as pd
import pickle
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import hf_hub_download
import requests
import os
import msgpack_numpy as m
import plotly.graph_objs as go
from sklearn.linear_model import LogisticRegression

torch.set_grad_enabled(False) # avoid blowing up mem

DEFAULT_EXAMPLE = text = "I really wished I could give this movie a higher rating. The plot was interesting, but the acting was terrible. The special effects were great, but the pacing was off. The movie was too long, but the ending was satisfying."

params = {
    "model_name" : "google/gemma-2-9b-it",
    "width" : "16k",
    "layer" : 31,
    "l0" : 76,
    "sae_repo_id": "google/gemma-scope-9b-it-res",
    "filename" : "layer_31/width_16k/average_l0_76/params.npz"
}

title = """
<div class='parent' align="center">
<div class='child' style="display: inline-block !important; margin-bottom: 20px;">
<h1 style="margin-bottom: 30px;">🔍Interpretable Classifier for movie ratings using Gemma 2 with SAEs</h1>
</div>
</div>
<div class='parent' align="center">
<p>This space demonstrates how a linear classifier trained on top of features learned by Sparse Auto Encoders (SAEs) can be used to create interpretable natural language classifiers.</p>
<p>We leverage the interpretability API of <b>Neuronpedia</b> to provide more information about the features used by the LLM (like what tokens activate it the most and their distribution).</p>
<p><b>More resources on interpretability for LLMs using SAEs:</b></p>
</div>
<ul>
    <li><a href="https://transformer-circuits.pub/2024/scaling-monosemanticity/">Anthropic: Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet</a></li>
    <li><a href="https://blog.eleuther.ai/autointerp/">EleutherAI: Open Source Automated Interpretability for Sparse Autoencoder Features</a></li>
    <li><a href="https://www.gemma.ai/gemma-scope">Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2</a></li>
</ul>
<div class='parent' align="center">
<p>About us: <b> 🌊 LaVague</b> is an open-source framework to build AI Web Agents. Check out our <a href="https://github.com/lavague-ai/LaVague">GitHub</a> or join our <a href="https://discord.com/invite/SDxn9KpqX9">Discord</a>.</p>
</div>
"""

css = """
    .my-button {
        height: 100px; /* Increase the height of the buttons */
        width: 100%; /* Make sure the button takes the full width */
        max-width: 300px; /* Optional: set a max width */
        max-height: 80px;
        font-size: 1.1rem; /* Increase font size */
    }
    .button-container {
        display: flex;
        justify-content: center; /* Center buttons horizontally */
        align-items: center; /* Center buttons vertically */
        height: 100%; /* Ensure it takes up the full height */
        width: 100%; /* Ensure it takes up the full width */
    }
"""

C = 0.01

model_name = params["model_name"]
width = params["width"]
layer = params["layer"]
l0 = params["l0"]
sae_repo_id = params["sae_repo_id"]
filename = params["filename"]

path_to_params = hf_hub_download(
    repo_id=sae_repo_id,
    filename=filename,
    force_download=False,
    token=os.environ['TOKEN'],
)

tokenizer =  AutoTokenizer.from_pretrained(model_name, token=os.environ['TOKEN'])

params = np.load(path_to_params)
pt_params = {k: torch.from_numpy(v) for k, v in params.items()}

clf_name = f"linear_classifier_C_{C}_ "+ model_name + "_" + filename.split(".npz")[0]
clf_name = clf_name.replace(os.sep, "_")

with open(f"{clf_name}.pkl", 'rb') as model_file:
    clf: LogisticRegression = pickle.load(model_file)


def get_feature_descriptions(feature):
    layer_name = f"{layer}-gemmascope-res-{width}"
    model_name_neuronpedia = model_name.split("/")[1]

    url = f"https://www.neuronpedia.org/api/feature/{model_name_neuronpedia}/{layer_name}/{feature}"

    response = requests.get(url)
    output = response.json()["explanations"][0]["description"]
    return output

def embed_content(url):
    html_content = f"""
    <div style="width:100%; height:500px; overflow:hidden;">
        <iframe src="{url}" width="100%" height="100%" frameborder="0"></iframe>
    </div>
    """
    return html_content

def dummy_function(*args):
    # This is a placeholder function. Replace with your actual logic.
    return "Scores will be displayed here"

examples = [
    "Despite moments of promise, this film ultimately falls short of its potential. The premise intrigues, offering a fresh take on a familiar genre, but the execution stumbles in crucial areas",
]

topk = 5

# Function to wrap in a FastAPI in case of
def get_activations(text):
    response = requests.post("http://34.71.249.22:3000/execute_req", json={"query": text})
    pack = m.unpackb(response.content)
    sae_act = torch.from_numpy(pack["sae_act"]).to(dtype=torch.bfloat16)
    return sae_act

def get_features(text):
    sae_act = get_activations(text)
    sae_act_aggregated = ((sae_act[:,:,:] > 0).sum(1) > 0).numpy()

    X = pd.DataFrame(sae_act_aggregated)

    feature_contributions = X.iloc[0].astype(float).values * clf.coef_[0]

    contrib_df = pd.DataFrame({
            'feature': range(len(feature_contributions)),
            'contribution': feature_contributions
    })

    contrib_df = contrib_df.loc[contrib_df['contribution'].abs() > 0]

    # Sort by absolute contribution and get top N
    contrib_df = contrib_df.reindex(contrib_df['contribution'].abs().sort_values(ascending=False).index)

    contrib_df = contrib_df.head(topk)
    descriptions = []
    for feature in contrib_df["feature"]:
        description = get_feature_descriptions(feature)
        print(description)
        descriptions.append(description)
    contrib_df["description"] = descriptions

    fig = go.Figure(go.Bar(
        x=contrib_df['contribution'],
        y=contrib_df['description'],
        orientation='h'  # Horizontal bar chart
    ))

    fig.update_layout(
        title='Feature contribution',
        xaxis_title='Contribution',
        yaxis_title='Features',
        height=500,
        margin=dict(l=200)  # Increase left margin to accommodate longer feature names
    )
    fig.update_yaxes(autorange="reversed")

    probability = clf.predict_proba(X)[0]
    classes = {
        "Positive": probability[1],
        "Negative": probability[0]
    }

    choices = [(description, feature) for description, feature in zip(contrib_df["description"], contrib_df["feature"])]
    dropdown = gr.Dropdown(choices=choices, 
                           value=choices[0][1],
                           interactive=True, label="Features")

    return classes, fig, dropdown

def get_highlighted_text(text, feature):
    inputs = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
    sae_act = get_activations(text)

    activated_tokens = sae_act[0:,:,feature]
    max_activation = activated_tokens.max().item()
    activated_tokens /= max_activation

    activated_tokens = activated_tokens.float().numpy()

    output = []

    for i, token_id in enumerate(inputs[0, :]):
        token = tokenizer.decode(token_id)
        output.append((token, activated_tokens[0, i]))

    return output

def get_feature_iframe(feature):
    layer_name = f"{layer}-gemmascope-res-{width}"
    model_name_neuronpedia = model_name.split("/")[1]
    url = f"https://neuronpedia.org/{model_name_neuronpedia}/{layer_name}/{feature}?embed=true"

    html_content = embed_content(url)
    html = gr.HTML(html_content)
    return html

with gr.Blocks(gr.themes.Default(primary_hue="blue", secondary_hue="neutral"), css=css) as demo:
    with gr.Tab(""):
        with gr.Row():
            gr.HTML(title)
        with gr.Row():
            with gr.Column(scale=4):
                input_text = gr.Textbox(label="Input", show_label=False, value=DEFAULT_EXAMPLE)
                gr.Examples(
                    examples=examples,
                    inputs=input_text,
                )
            with gr.Column(scale=1):
                run_button = gr.Button("Run")

        with gr.Row():
            label = gr.Label(label="Scores")

        with gr.Row():
            with gr.Column(scale=1):
                plot = gr.Plot(label="Plot")
                dropdown = gr.Dropdown(choices=["Option 1"], label="Features")
            
            with gr.Column(scale=1):
                highlighted_text = gr.HighlightedText(
                                    label="Activating Tokens",
                                    combine_adjacent=True,
                                    show_legend=True,
                                    color_map={"+": "red", "-": "green"})

        with gr.Row():
            html = gr.HTML()

        # Connect the components
        run_button.click(
            fn=get_features,
            inputs=[input_text],
            outputs=[label, plot, dropdown],
        ).then(
        fn=get_highlighted_text,
            inputs=[input_text, dropdown],
            outputs=[highlighted_text]
        ).then(
        fn=get_feature_iframe,
            inputs=[dropdown],
            outputs=[html]
        )

        dropdown.change(
            fn=get_highlighted_text,
            inputs=[input_text, dropdown],
            outputs=[highlighted_text]
        ).then(
        fn=get_feature_iframe,
            inputs=[dropdown],
            outputs=[html]
        )

demo.launch(share=True)