Update app.py

app.py

@@ -45,25 +45,22 @@ suggested_interpretation_prompts = ["Before responding, let me repeat the messag
 def initialize_gpu():
     pass
 
-def get_hidden_states(raw_original_prompt, interpretation_args, interpretation_outputs):
+def get_hidden_states(raw_original_prompt):
     original_prompt = original_prompt_template.format(prompt=raw_original_prompt)
     model_inputs = tokenizer(original_prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
     tokens = tokenizer.batch_decode(model_inputs.input_ids[0])
     outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
     hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
-    token_btns = []
-    for i, token in enumerate(tokens):
-        btn = gr.Button(token)
-        btn.click(partial(run_interpretation, interpreted_vectors=hidden_states[:, i]),
-                  interpretation_args, interpretation_outputs)
-        token_btns.append(btn)
-    token_btns += [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))]
-    return token_btns
-
-
-def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample, 
-              temperature, top_k, top_p, repetition_penalty, length_penalty, interpreted_vectors, num_beams=1):
-    
+    token_btns = ([gr.Button(token, visible=True) for token in tokens] 
+                  + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
+    return [hidden_state, *token_btns]
+
+
+def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
+                       temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
+                       num_beams=1):
+
+    interpreted_vectors = global_state[:, i]
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
     # generation parameters
@@ -83,7 +80,7 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
     interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
     # generate the interpretations
-    generated = interpretation_prompt.generate(model, {0: hidden_states[:, -1]}, k=3, **generation_kwargs)
+    generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     return generation_texts
 
@@ -105,6 +102,8 @@ model = AutoModelClass.from_pretrained(model_name, **model_args)
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, token=os.environ['hf_token'])
 
 # demo
+global_state = gr.State([])
+json_output = gr.JSON()
 with gr.Blocks(theme=gr.themes.Default()) as demo:
     with gr.Row():
         with gr.Column(scale=5):
@@ -144,15 +143,15 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
 
     with gr.Group('Output'):
+        tokens_container = []
         with gr.Row():
-            tokens_container = [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS)]
-        json_output = gr.JSON()
-
-    interpretation_args = [interpretation_prompt, num_tokens, do_sample, temperature, 
-                           top_k, top_p, repetition_penalty, length_penalty]
-    interpretation_outputs = [json_output]
-
-    original_prompt_btn.click(partial(get_hidden_states, interpretation_args=interpretation_args,
-                                      interpretation_outputs=interpretation_outputs
-                                     ), [original_prompt_raw], [*tokens_container])    
+            for _ in range(MAX_PROMPT_TOKENS):
+                btn = gr.Button('', visible=False)
+                btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, num_tokens, do_sample, temperature, 
+                                                           top_k, top_p, repetition_penalty, length_penalty
+                                                          ], [json_output])
+            tokens_container.append(btn)
+        json_output.render()
+        
+    original_prompt_btn.click(get_hidden_states, [original_prompt_raw], [global_state, *tokens_container])    
     demo.launch()