app.py

import gradio as gr
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import hf_hub_download
import numpy as np
import torch
import pickle
import numpy as np
import pandas as pd
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"
}

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

C = 0.01

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='auto',
    torch_dtype=torch.bfloat16,
)
tokenizer =  AutoTokenizer.from_pretrained(model_name)

path_to_params = hf_hub_download(
    repo_id=sae_repo_id,
    filename=filename,
    force_download=False,
)

params = np.load(path_to_params)
pt_params = {k: torch.from_numpy(v).cuda() 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)

import torch.nn as nn
class JumpReLUSAE(nn.Module):
  def __init__(self, d_model, d_sae):
    # Note that we initialise these to zeros because we're loading in pre-trained weights.
    # If you want to train your own SAEs then we recommend using blah
    super().__init__()
    self.W_enc = nn.Parameter(torch.zeros(d_model, d_sae))
    self.W_dec = nn.Parameter(torch.zeros(d_sae, d_model))
    self.threshold = nn.Parameter(torch.zeros(d_sae))
    self.b_enc = nn.Parameter(torch.zeros(d_sae))
    self.b_dec = nn.Parameter(torch.zeros(d_model))

  def encode(self, input_acts):
    pre_acts = input_acts @ self.W_enc + self.b_enc
    mask = (pre_acts > self.threshold)
    acts = mask * torch.nn.functional.relu(pre_acts)
    return acts

  def decode(self, acts):
    return acts @ self.W_dec + self.b_dec

  def forward(self, acts):
    acts = self.encode(acts)
    recon = self.decode(acts)
    return recon

sae = JumpReLUSAE(params['W_enc'].shape[0], params['W_enc'].shape[1])
sae.load_state_dict(pt_params)
sae.to(dtype=torch.bfloat16).cuda()

@torch.no_grad()
def gather_residual_activations(model, target_layer, inputs):
  target_act = None
  def gather_target_act_hook(mod, inputs, outputs):
    nonlocal target_act # make sure we can modify the target_act from the outer scope
    target_act = outputs[0]
    return outputs
  handle = model.model.layers[target_layer].register_forward_hook(gather_target_act_hook)
  _ = model.forward(inputs)
  handle.remove()
  return target_act

import requests

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

def get_features(text):

    inputs = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True).to("cuda")

    target_act = gather_residual_activations(model, layer, inputs)
    sae_act = sae.encode(target_act)
    sae_act_aggregated = ((sae_act[:,:,:] > 0).sum(1) > 0).cpu().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

    import plotly.graph_objs as go

    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).to("cuda")

    target_act = gather_residual_activations(model, layer, inputs)
    sae_act = sae.encode(target_act)

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

    activated_tokens = activated_tokens.cpu().detach().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)