add our gemma residuals

app.py

@@ -7,10 +7,8 @@ import os
 import numpy as np
 
 
-token_table = pq.read_table("weights/tokens.parquet")
 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")
 
@@ -20,7 +18,7 @@ def find_revions():
         if parquet.endswith(".parquet"):
             parts = parquet.split("-")
             if len(parts) > 2:
-                revisions.add(int(parts[2][1:]))
+                revisions.add(int(parts[-3][1:]))
     return sorted(revisions)
 
 def find_layers(revision):
@@ -28,8 +26,8 @@ def find_layers(revision):
     for parquet in parquets:
         if parquet.endswith(".parquet"):
             parts = parquet.split("-")
-            if len(parts) > 2 and int(parts[2][1:]) == revision:
-                layers.add(int(parts[1][1:]))
+            if len(parts) > 2 and int(parts[-3][1:]) == revision:
+                layers.add(int(parts[-4][1:]))
     return sorted(layers)
 
 revisions = find_revions()
@@ -51,13 +49,29 @@ nearby = 8
 stride = 0.25
 n_bins = 10
 
-def make_cache_name(layer, revision):
-    return f"{cache_path}/phi-l{layer}-r{revision}-st0.25x128-activations.parquet"
+def make_cache_name(layer, revision, model):
+    return f"{cache_path}/{model}-l{layer}-r{revision}-st0.25x128-activations.parquet"
+
+models = {
+    "gemma-2b-r": "gemma-2b-residuals",
+    "phi-3": "phi"
+}
+
+tokenizers = {
+    "gemma-2b-r": "alpindale/gemma-2b",
+    "phi-3": "microsoft/Phi-3-mini-4k-instruct"
+}
+
+token_tables = {
+    "gemma-2b-r": pq.read_table("weights/tokens_gemma.parquet"),
+    "phi-3": pq.read_table("weights/tokens.parquet")
+}
 
 with gr.Blocks() as demo:
     feature_table = gr.State(None)
 
-    tokenizer_name = gr.Textbox(TOKENIZER, label="Tokenizer")
+    model_name = gr.Dropdown(["phi-3", "gemma-2b-r"], label="Model")
+
     revision_dropdown = gr.Dropdown(revisions, label="Revision")
 
     layer_dropdown = gr.Dropdown(layers[4], label="Layer")
@@ -86,9 +100,17 @@ with gr.Blocks() as demo:
     cm = gr.HighlightedText()
     frame = gr.Highlightedtext()
 
-    def update(revision, layer, feature, extra_tokens, tokenizer_name):
+    def update(model, revision, layer, feature, extra_tokens):
+        correction = 1
+        if "gemma" in model:
+            correction = 0
+
+        token_table = token_tables[model]
+
+        tokenizer_name = tokenizers[model]
+
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-        table = pq.read_table(make_cache_name(layer, revision))
+        table = pq.read_table(make_cache_name(layer, revision, models[model]))
         table_feat = table.filter(pc.field("feature") == feature).to_pandas()
 
         # freq_t = table_feat[["activation", "freq"]]
@@ -100,7 +122,7 @@ with gr.Blocks() as demo:
         table_feat = table_feat.sort_values("activation", ascending=False)
 
         texts = table_feat["token"].apply(
-            lambda x: [tokenizer.decode(y).replace("\n", " ") for y in token_table[max(0, x - nearby + 1 - extra_tokens):x + extra_tokens + nearby + 2]["tokens"].to_numpy()]
+            lambda x: [tokenizer.decode(y).replace("\n", " ") for y in token_table[max(0, x - nearby + correction - extra_tokens):x + extra_tokens + nearby + 1 + correction]["tokens"].to_numpy()]
         ).tolist()
 
         # texts = [tokenizer.tokenize(text) for text in texts]
@@ -145,7 +167,7 @@ with gr.Blocks() as demo:
         
 
     # feature_dropdown.change(update, [layer_dropdown, feature_dropdown, tokenizer_name], [frame, cm, frequency, autoi_expl, selfe_expl])
-    feature_input.change(update, [revision_dropdown, layer_dropdown, feature_input, extra_tokens,   tokenizer_name], [frame, cm, frequency, autoi_expl, selfe_expl])
+    feature_input.change(update, [model_name, revision_dropdown, layer_dropdown, feature_input, extra_tokens], [frame, cm, frequency, autoi_expl, selfe_expl])
 
 
 if __name__ == "__main__":

weights/caches/gemma-2b-residuals-l10-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l11-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l12-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l13-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l14-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l15-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l16-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l6-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l8-r1-st0.25x128-activations.parquet

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

weights/caches/gemma-2b-residuals-l9-r1-st0.25x128-activations.parquet

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

weights/tokens_gemma.parquet

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