Update app.py

app.py

@@ -56,9 +56,9 @@ def get_hidden_states(progress, raw_original_prompt):
     return [hidden_states, *token_btns]
 
 
-def run_interpretation(progress, global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
+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):
+                       num_beams=1, progress=gr.Progress()):
 
     interpreted_vectors = global_state[:, i]
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
@@ -184,10 +184,9 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
                 tokens_container.append(btn)
             interpretation_bubbles = [gr.Textbox('', container=False, visible=False, elem_classes=['bubble'])
                                      for i in range(model.config.num_hidden_layers)]
-            progress = gr.Progress()
+            
         for i, btn in enumerate(tokens_container):
-            btn.click(partial(run_interpretation, i=i), [progress, 
-                                                         global_state, interpretation_prompt, num_tokens, do_sample, temperature, 
+            btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, num_tokens, do_sample, temperature, 
                                                          top_k, top_p, repetition_penalty, length_penalty
                                                               ], [*interpretation_bubbles])