Update app.py

app.py

@@ -74,7 +74,7 @@ def initialize_gpu():
     pass
 
 
-def reset_model(model_name, return_state=False): 
+def reset_model(model_name): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -90,8 +90,6 @@ def reset_model(model_name, return_state=False):
     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()
-    if return_state:
-        return global_state
     
 
 def get_hidden_states(raw_original_prompt):
@@ -145,11 +143,13 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
 
 ## main
 torch.set_grad_enabled(False)
+global_state = GlobalState()
 
 model_name = 'LLAMA2-7B'
-global_state = reset_model(model_name, return_state=True)
+reset_model(model_name)
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []
+
 for i in range(MAX_PROMPT_TOKENS):
     btn = gr.Button('', visible=False, elem_classes=['token_btn'])
     tokens_container.append(btn)