Update app.py

app.py

@@ -60,7 +60,7 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
     AutoModelClass = CAutoModelForCausalLM if use_ctransformers else AutoModelForCausalLM
     
     # get model
-    global_state.model, global_state.tokenizer, global_state.local_state = None, None, LocalState()
+    global_state.model, global_state.tokenizer, global_state.local_state.hidden_states = None, None, None
     gc.collect()
     global_state.model = AutoModelClass.from_pretrained(model_path, **model_args)
     if not dont_cuda:
@@ -227,7 +227,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                                                     ], [progress_dummy, *interpretation_bubbles])
     
     original_prompt_btn.click(get_hidden_states, 
-                              [gr.State(global_state), original_prompt_raw], 
+                              [gr.State(global_state.local_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)