Update app.py

app.py

@@ -27,6 +27,7 @@ class GlobalState:
     tokenizer : Optional[PreTrainedTokenizer] = None
     model : Optional[PreTrainedModel] = None
     local_state : LocalState = LocalState()
+    wait_with_hidden_state : bool = False
     interpretation_prompt_template : str = '{prompt}'
     original_prompt_template : str = 'User: [X]\n\nAnswer: {prompt}'
     layers_format : str = 'model.layers.{k}'
@@ -48,7 +49,6 @@ def initialize_gpu():
 
 def reset_model(model_name, *extra_components, with_extra_components=True): 
     # extract model info
-    
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
     global_state.original_prompt_template = model_args.pop('original_prompt_template')
@@ -57,6 +57,7 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
     tokenizer_path = model_args.pop('tokenizer') if 'tokenizer' in model_args else model_path
     use_ctransformers = model_args.pop('ctransformers', False)
     dont_cuda = model_args.pop('dont_cuda', False)
+    global_state.wait_with_hidden_states = model_args.pop('wait_with_hidden_states', False)
     AutoModelClass = CAutoModelForCausalLM if use_ctransformers else AutoModelForCausalLM
     
     # get model
@@ -80,22 +81,28 @@ def get_hidden_states(raw_original_prompt):
     model_inputs = tokenizer(original_prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
     tokens = tokenizer.batch_decode(model_inputs.input_ids[0])
     outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
-    hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
+    if global_state.wait_with_hidden_states:
+        global_state.local_state.hidden_states = None
+    else:
+        hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
+        global_state.local_state.hidden_states = hidden_states.cpu().detach()
+        
     token_btns = ([gr.Button(token, visible=True) for token in tokens] 
                   + [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.local_state.hidden_states = hidden_states.cpu().detach()
     return [progress_dummy_output, *token_btns, *invisible_bubbles]
 
 
 @spaces.GPU
-def run_interpretation(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, i, 
                        num_beams=1):
     model = global_state.model
     tokenizer = global_state.tokenizer
     print(f'run {model}')
+    if global_state.wait_with_hidden_states and global_state.local_state.hidden_states is None:
+        get_hidden_states(raw_original_prompt)
     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
 
@@ -218,7 +225,7 @@ 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), [interpretation_prompt, 
+        btn.click(partial(run_interpretation, i=i), [original_prompt_raw, interpretation_prompt, 
                                                      num_tokens, do_sample, temperature, 
                                                      top_k, top_p, repetition_penalty, length_penalty
                                                     ], [progress_dummy, *interpretation_bubbles])