Update app.py

app.py

@@ -68,9 +68,10 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
                 + [gr.Textbox('', visible=False) for _ in range(len(interpretation_bubbles))] 
                 + [gr.Button('', visible=False) for _ in range(len(tokens_container))]
                 + [*extra_components])
-    
+
+
 @spaces.GPU
-def get_hidden_states(raw_original_prompt):
+def get_hidden_states(global_state, raw_original_prompt):
     model, tokenizer = global_state.model, global_state.tokenizer
     original_prompt = global_state.original_prompt_template.format(prompt=raw_original_prompt)
     model_inputs = tokenizer(original_prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
@@ -222,7 +223,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                                                     ], [progress_dummy, *interpretation_bubbles])
     
     original_prompt_btn.click(get_hidden_states, 
-                              [original_prompt_raw], 
+                              [gr.State(global_state), original_prompt_raw], 
                               [progress_dummy, *tokens_container, *interpretation_bubbles])
     original_prompt_raw.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)