Update app.py

app.py

@@ -50,14 +50,19 @@ def get_hidden_states(raw_original_prompt):
     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)
-    # with gr.Row() as tokens_container:
-        # for token in tokens:
-        #     gr.Button(token)
-    return [gr.Button(tokens[i], visible=True) if i < len(tokens) else gr.Button('', visible=False) for i in range(MAX_PROMPT_TOKENS)]
-
-
-def run_model(raw_original_prompt, raw_interpretation_prompt, max_new_tokens, do_sample, 
-              temperature, top_k, top_p, repetition_penalty, length_penalty, num_beams=1):
+    token_btns = []
+    for i, token in enumerate(tokens):
+        btn = gr.Button(token)
+        btn.click(partial(run_interpretation, interpreted_vectors=hidden_states[:, i]),
+                  [interpretation_prompt, num_tokens, do_sample, temperature, top_k, top_p, repetition_penalty, length_penalty],
+                 [json_output])
+        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):
     
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
@@ -77,16 +82,9 @@ def run_model(raw_original_prompt, raw_interpretation_prompt, max_new_tokens, do
     interpretation_prompt = interpretation_prompt_template.format(prompt=raw_interpretation_prompt)
     interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
-    # compute the hidden stated from the original prompt (after putting it in the right template)
-    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)
-    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)
-
     # generate the interpretations
     generated = interpretation_prompt.generate(model, {0: hidden_states[:, -1]}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
-    # tokens = [x.lstrip('▁') for x in tokenizer.tokenize(text)]
     return generation_texts
 
 
@@ -148,13 +146,7 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
     with gr.Group('Output'):
         with gr.Row():
             tokens_container = [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS)]
-        with gr.Column() as interpretations_container:
-            pass
-
-    original_prompt_btn.click(get_hidden_states, [original_prompt_raw], [*tokens_container])
-    # btn.click(run_model, 
-    #           [text, interpretation_prompt, num_tokens, do_sample, temperature, 
-    #            top_k, top_p, repetition_penalty, length_penalty], 
-    #           [tokens_container])
-    
+        json_output = gr.JSON()
+        
+    original_prompt_btn.click(get_hidden_states, [original_prompt_raw], [*tokens_container])    
     demo.launch()