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(progress, raw_original_prompt):
+def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
     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])
@@ -191,6 +191,6 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
                                                               ], [*interpretation_bubbles])
         
     original_prompt_btn.click(get_hidden_states, 
-                              [progress, original_prompt_raw], 
+                              [original_prompt_raw], 
                               [global_state, *tokens_container])    
     demo.launch()