Update app.py

app.py

@@ -18,11 +18,15 @@ MAX_PROMPT_TOKENS = 60
 MAX_NUM_LAYERS = 50
 welcome_message = '**You are now running {model_name}!!** 🥳🥳🥳'
 
+@dataclass
+class LocalState:
+    hidden_states: Optional[torch.Tensor] = None
+
 @dataclass
 class GlobalState:
     tokenizer : Optional[PreTrainedTokenizer] = None
     model : Optional[PreTrainedModel] = None
-    hidden_states : Optional[torch.Tensor] = None
+    local_state : LocalState = LocalState()
     interpretation_prompt_template : str = '{prompt}'
     original_prompt_template : str = 'User: [X]\n\nAnswer: {prompt}'
     layers_format : str = 'model.layers.{k}'
@@ -56,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.hidden_states = None, None, None
+    global_state.model, global_state.tokenizer, global_state.local_state = None, None, LocalState()
     gc.collect()
     global_state.model = AutoModelClass.from_pretrained(model_path, **model_args)
     if not dont_cuda:
@@ -71,7 +75,7 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
 
 
 @spaces.GPU
-def get_hidden_states(global_state, raw_original_prompt):
+def get_hidden_states(local_state, 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)
@@ -82,7 +86,7 @@ def get_hidden_states(global_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)]
-    global_state.hidden_states = hidden_states.cpu().detach().numpy()
+    local_state.hidden_states = hidden_states.cpu().detach()
     return [progress_dummy_output, *token_btns, *invisible_bubbles]
 
 
@@ -93,7 +97,7 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
     model = global_state.model
     tokenizer = global_state.tokenizer
     print(f'run {model}')
-    interpreted_vectors = torch.tensor(global_state.hidden_states[:, i]).to(model.device).to(model.dtype)
+    interpreted_vectors = torch.tensor(global_state.local_state.hidden_states[:, i]).to(model.device).to(model.dtype)
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
     # generation parameters