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

🔍Interpretable Classifier for movie ratings using Gemma 2 with SAEs

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.

We leverage the interpretability API of Neuronpedia to provide more information about the features used by the LLM (like what tokens activate it the most and their distribution).

More resources on interpretability for LLMs using SAEs:

About us: 🌊 LaVague is an open-source framework to build AI Web Agents. Check out our GitHub or join our Discord.

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