Update app.py

app.py

@@ -87,7 +87,7 @@ def get_hidden_states(local_state, raw_original_prompt):
     progress_dummy_output = ''
     invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_NUM_LAYERS)]
     local_state.hidden_states = hidden_states.cpu().detach()
-    return [progress_dummy_output, *token_btns, *invisible_bubbles]
+    return [progress_dummy_output, local_state, *token_btns, *invisible_bubbles]
 
 
 @spaces.GPU
@@ -217,18 +217,19 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                                              elem_classes=['bubble', 'even_bubble' if i % 2 == 0 else 'odd_bubble']
                                             ) for i in range(MAX_NUM_LAYERS)]
     
+
+    local_state = gr.State(global_state.local_state)
     
-    # event listeners
-    
+    # event listeners    
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [interpretation_prompt, 
                                                      num_tokens, do_sample, temperature, 
                                                      top_k, top_p, repetition_penalty, length_penalty
                                                     ], [progress_dummy, *interpretation_bubbles])
-    
+
     original_prompt_btn.click(get_hidden_states, 
-                              [gr.State(global_state.local_state), original_prompt_raw], 
-                              [progress_dummy, *tokens_container, *interpretation_bubbles])
+                              [local_state, original_prompt_raw], 
+                              [progress_dummy, local_state, *tokens_container, *interpretation_bubbles])
     original_prompt_raw.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)
 
     extra_components = [interpretation_prompt, original_prompt_raw, original_prompt_btn]