Update app.py

app.py

@@ -57,7 +57,7 @@ suggested_interpretation_prompts = [
 def initialize_gpu():
     pass
 
-def reset_model(global_state, model_name, load_on_gpu, *extra_components, reset_sentence_transformer=False, with_extra_components=True): 
+def reset_model(model_name, load_on_gpu, *extra_components, reset_sentence_transformer=False, with_extra_components=True): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -84,15 +84,15 @@ def reset_model(global_state, model_name, load_on_gpu, *extra_components, reset_
     global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
     gc.collect()
     if with_extra_components:
-        return ([global_state, welcome_message.format(model_name=model_name)] 
+        return ([welcome_message.format(model_name=model_name)] 
                 + [gr.Textbox('', visible=False) for _ in range(len(interpretation_bubbles))] 
                 + [gr.Button('', visible=False) for _ in range(len(tokens_container))]
                 + [*extra_components])
     else:
-        return global_state
+        return None
 
 
-def get_hidden_states(global_state, raw_original_prompt, force_hidden_states=False):
+def get_hidden_states(raw_original_prompt, force_hidden_states=False):
     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)
@@ -120,7 +120,7 @@ def get_hidden_states(global_state, raw_original_prompt, force_hidden_states=Fal
 
 
 @spaces.GPU
-def run_interpretation(global_state, raw_original_prompt, raw_interpretation_prompt, max_new_tokens, do_sample, 
+def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_tokens, do_sample, 
                        temperature, top_k, top_p, repetition_penalty, length_penalty, use_gpu, i, 
                        num_beams=1):
     model = global_state.model
@@ -197,8 +197,8 @@ for i in range(MAX_PROMPT_TOKENS):
     tokens_container.append(btn)
 
 with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
-    global_state = gr.State(partial(reset_model, GlobalState(), 
-                            model_name, load_on_gpu=True, with_extra_components=False, reset_sentence_transformer=True))
+    global_state = GlobalState()
+    reset_model(model_name, load_on_gpu=True, with_extra_components=False, reset_sentence_transformer=True)
     with gr.Row():
         with gr.Column(scale=5):
             gr.Markdown('# 😎 Self-Interpreting Models')
@@ -278,19 +278,19 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
         
     # event listeners    
     for i, btn in enumerate(tokens_container):
-        btn.click(partial(run_interpretation, i=i), [global_state, raw_original_prompt, raw_interpretation_prompt, 
+        btn.click(partial(run_interpretation, i=i), [raw_original_prompt, raw_interpretation_prompt, 
                                                      num_tokens, do_sample, temperature, 
                                                      top_k, top_p, repetition_penalty, length_penalty,
                                                      use_gpu
                                                     ], [progress_dummy, *interpretation_bubbles])
 
     original_prompt_btn.click(get_hidden_states, 
-                              [global_state, raw_original_prompt], 
+                              [raw_original_prompt], 
                               [progress_dummy, *tokens_container, *interpretation_bubbles])
     raw_original_prompt.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)
 
     extra_components = [raw_interpretation_prompt, raw_original_prompt, original_prompt_btn]
-    model_chooser.change(reset_model, [global_state, model_chooser, load_on_gpu, *extra_components], 
-                         [global_state, welcome_model, *interpretation_bubbles, *tokens_container, *extra_components])
+    model_chooser.change(reset_model, [model_chooser, load_on_gpu, *extra_components], 
+                         [welcome_model, *interpretation_bubbles, *tokens_container, *extra_components])
     
     demo.launch()