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"
}

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)

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() as demo:
    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)