Update app.py

app.py

@@ -74,8 +74,7 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
                 + [*extra_components])
 
 
-@spaces.GPU
-def get_hidden_states(local_state, raw_original_prompt):
+def get_hidden_states(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)
@@ -86,8 +85,8 @@ def get_hidden_states(local_state, raw_original_prompt):
                   + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
     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, local_state, *token_btns, *invisible_bubbles]
+    global_state.local_state.hidden_states = hidden_states.cpu().detach()
+    return [progress_dummy_output, *token_btns, *invisible_bubbles]
 
 
 @spaces.GPU
@@ -216,10 +215,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
         interpretation_bubbles = [gr.Textbox('', container=False, visible=False, 
                                              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    
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [interpretation_prompt, 
@@ -228,8 +224,8 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                                                     ], [progress_dummy, *interpretation_bubbles])
 
     original_prompt_btn.click(get_hidden_states, 
-                              [local_state, original_prompt_raw], 
-                              [progress_dummy, local_state, *tokens_container, *interpretation_bubbles])
+                              [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)
 
     extra_components = [interpretation_prompt, original_prompt_raw, original_prompt_btn]