Add explanations

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import pyarrow.parquet as pq
 import pyarrow.compute as pc
 from transformers import AutoTokenizer
+from datasets import load_dataset
 import os
 import numpy as np
 
@@ -11,36 +12,50 @@ cache_path = "weights/caches"
 parquets = os.listdir(cache_path)
 TOKENIZER = "microsoft/Phi-3-mini-4k-instruct"
 
+dataset = load_dataset("kisate-team/feature-explanations", split="train")
+
+layers = dataset.unique("layer")
+
+features = {layer:{item["feature"]:item for item in dataset if item["layer"] == layer} for layer in layers}
+
 nearby = 8
 stride = 0.25
 n_bins = 10
 
+def make_cache_name(layer):
+    return f"{cache_path}/phi-l{layer}-r4-st0.25x128-activations.parquet"
+
 with gr.Blocks() as demo:
     feature_table = gr.State(None)
 
     tokenizer_name = gr.Textbox(TOKENIZER)
-    dropdown = gr.Dropdown(parquets)
-    feature_input = gr.Number(0)
-    token_range = gr.Number(64)
+    layer_dropdown = gr.Dropdown(layers)
+    feature_dropdown = gr.Dropdown()
+
+    def update_features(layer):
+        feature_dropdown = gr.Dropdown(features[layer].keys())
+        return feature_dropdown
+
+    layer_dropdown.input(update_features, layer_dropdown, feature_dropdown)
+
 
     frequency = gr.Number(0, label="Total frequency (%)")
-    histogram = gr.LinePlot(x="activation", y="freq")
+    # histogram = gr.LinePlot(x="activation", y="freq")
+
+    autoi_expl = gr.Textbox()
+    selfe_expl = gr.Textbox()
+
     cm = gr.HighlightedText()
-    frame = gr.Highlightedtext(
-        show_legend=True
-    )
+    frame = gr.Highlightedtext()
 
-    def update(cache_name, feature, tokenizer_name, token_range):
-        if cache_name is None:
-            return
+    def update(layer, feature, tokenizer_name):
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-        table = pq.read_table(f"{cache_path}/{cache_name}")
+        table = pq.read_table(make_cache_name(layer))
         table_feat = table.filter(pc.field("feature") == feature).to_pandas()
 
-        freq_t = table_feat[["activation", "freq"]]
+        # freq_t = table_feat[["activation", "freq"]]
         total_freq = float(table_feat["freq"].sum()) * 100
         
-
         table_feat = table_feat[table_feat["activation"] > 0]
         table_feat = table_feat[table_feat["freq"] > 0]
 
@@ -70,11 +85,19 @@ with gr.Blocks() as demo:
             flat_data = []
             color_map_data = []
 
-        return flat_data, color_map_data, freq_t, total_freq
+        autoi_expl = features[layer][feature]["explanation"]
+        selfe_expl = features[layer][feature]["gen_explanations"]
+
+        if selfe_expl is not None:
+            selfe_expl = "\n".join(
+                f"{i+1}. \"{x}\"" for i, x in enumerate(selfe_expl)
+            )
+
+        return flat_data, color_map_data, total_freq, autoi_expl, selfe_expl
         
 
-    dropdown.change(update, [dropdown, feature_input, tokenizer_name, token_range], [frame, cm, histogram, frequency])
-    feature_input.change(update, [dropdown, feature_input, tokenizer_name, token_range], [frame, cm, histogram, frequency])
+    feature_dropdown.change(update, [layer_dropdown, feature_dropdown, tokenizer_name], [frame, cm, frequency, autoi_expl, selfe_expl])
+    # feature_input.change(update, [dropdown, feature_input, tokenizer_name, token_range], [frame, cm, histogram, frequency])
 
 
 if __name__ == "__main__":