Update app.py

app.py

@@ -53,12 +53,13 @@ def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
     hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
     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_states, *token_btns]
+    progress_dummy_output = ''
+    return [progress_dummy_output, hidden_states, *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, progress=gr.Progress()):
+                       num_beams=1):
 
     interpreted_vectors = global_state[:, i]
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
@@ -82,7 +83,8 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     # generate the interpretations
     generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
-    return [gr.Textbox(text, visible=True, container=False) for text in generation_texts]
+    progress_dummy_output = ''
+    return [progress_dummy_output] + [gr.Textbox(text, visible=True, container=False) for text in generation_texts]
 
 
 ## main
@@ -182,15 +184,18 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             for i in range(MAX_PROMPT_TOKENS):
                 btn = gr.Button('', visible=False, elem_classes=['token_btn'])
                 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_dummy = gr.Text('', container=False)
+        interpretation_bubbles = [gr.Textbox('', container=False, visible=False, elem_classes=['bubble'])
+                                 for i in range(model.config.num_hidden_layers)]
+        
         for i, btn in enumerate(tokens_container):
-            btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, num_tokens, do_sample, temperature, 
+            btn.click(partial(run_interpretation, i=i), [progress_dummy, global_state, interpretation_prompt, 
+                                                         num_tokens, do_sample, temperature, 
                                                          top_k, top_p, repetition_penalty, length_penalty
                                                               ], [*interpretation_bubbles])
         
     original_prompt_btn.click(get_hidden_states, 
                               [original_prompt_raw], 
-                              [global_state, *tokens_container])    
+                              [progress_dummy, global_state, *tokens_container])    
     demo.launch()