extra tokens option

app.py

@@ -71,6 +71,7 @@ with gr.Blocks() as demo:
         return layer_dropdown
 
     frequency = gr.Number(0, label="Total frequency (%)")
+    extra_tokens = gr.Number(0, label="Extra Max Act Tokens")
 
     # layer_dropdown.input(update_features, layer_dropdown, feature_dropdown)
     # histogram = gr.LinePlot(x="activation", y="freq")
@@ -85,7 +86,7 @@ with gr.Blocks() as demo:
     cm = gr.HighlightedText()
     frame = gr.Highlightedtext()
 
-    def update(revision, layer, feature, tokenizer_name):
+    def update(revision, layer, feature, extra_tokens, tokenizer_name):
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
         table = pq.read_table(make_cache_name(layer, revision))
         table_feat = table.filter(pc.field("feature") == feature).to_pandas()
@@ -99,13 +100,13 @@ 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):x + nearby + 2]["tokens"].to_numpy()]
+            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()]
         ).tolist()
 
         # texts = [tokenizer.tokenize(text) for text in texts]
         activations = table_feat["nearby"].to_numpy()
 
-        activations = [a for i, a in enumerate(activations) if len(texts[i]) > 0]
+        activations = [[0] * extra_tokens + a.tolist()  + [0] * extra_tokens for i, a in enumerate(activations) if len(texts[i]) > 0]
         texts = [text for text in texts if len(text) > 0]
 
         for t, a in zip(texts, activations):
@@ -144,7 +145,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, 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])
 
 
 if __name__ == "__main__":