Update app.py

app.py

@@ -45,7 +45,7 @@ suggested_interpretation_prompts = ["Before responding, let me repeat the messag
 def initialize_gpu():
     pass
 
-def get_hidden_states(raw_original_prompt):
+def get_hidden_states(raw_original_prompt, interpretation_args, interpretation_outputs):
     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])
@@ -55,8 +55,7 @@ def get_hidden_states(raw_original_prompt):
     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])
+                  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
@@ -143,11 +142,17 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
 
     with gr.Group('Interpretation'):
         interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
-    
+
     with gr.Group('Output'):
         with gr.Row():
             tokens_container = [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS)]
         json_output = gr.JSON()
-        
-    original_prompt_btn.click(get_hidden_states, [original_prompt_raw], [*tokens_container])    
+
+    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])    
     demo.launch()