Upload 2 files

app.py

@@ -0,0 +1,272 @@
+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)

linear_classifier_C_0.01_ google_gemma-2-9b-it_layer_31_width_16k_average_l0_76_params.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:facb7f240c7957bd7c2ebf5d4f472e100eefe20c2a498e0fbc35e25754bf9283
+size 131805