Update app.py

app.py

@@ -74,7 +74,7 @@ def initialize_gpu():
     pass
 
 
-def reset_model(model_name, global_state):
+def reset_model(model_name, return_state=False): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -90,10 +90,11 @@ def reset_model(model_name, global_state):
     global_state.model = AutoModelClass.from_pretrained(model_path, **model_args).cuda()
     global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
     gc.collect()
-    return global_state
+    if return_state:
+        return global_state
     
 
-def get_hidden_states(global_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)
@@ -105,11 +106,11 @@ def get_hidden_states(global_state, raw_original_prompt):
     progress_dummy_output = ''
     invisible_bubbles = [gr.Textbox('', visible=False) for i in range(len(interpretation_bubbles))]
     global_state.hidden_states = hidden_states
-    return [progress_dummy_output, global_state, *token_btns, *invisible_bubbles]
+    return [progress_dummy_output, *token_btns, *invisible_bubbles]
 
 
 @spaces.GPU
-def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
+def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample, 
                        temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
                        num_beams=1):
 
@@ -143,7 +144,7 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
 
 
 ## main
-
+global_state = reset_model(model_name, return_state=True)
 torch.set_grad_enabled(False)
 model_name = 'LLAMA2-7B'
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
@@ -153,7 +154,6 @@ 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(reset_model(model_name, GlobalState()))
     
     with gr.Row():
         with gr.Column(scale=5):
@@ -236,8 +236,9 @@ 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(model.config.num_hidden_layers)]
 
+    
     # event listeners
-    model_chooser.change(reset_new_model, [model_chooser, global_state], [global_state])
+    model_chooser.change(reset_new_model, [model_chooser], [])
     
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, 
@@ -247,6 +248,6 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     
     original_prompt_btn.click(get_hidden_states, 
                               [original_prompt_raw], 
-                              [progress_dummy, global_state, *tokens_container, *interpretation_bubbles])
+                              [progress_dummy, *tokens_container, *interpretation_bubbles])
     original_prompt_raw.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)
     demo.launch()