Update app.py

app.py

@@ -75,15 +75,15 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
                 + [*extra_components])
 
 
-def get_hidden_states(raw_original_prompt):
+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)
     tokens = tokenizer.batch_decode(model_inputs.input_ids[0])
-    outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
-    if global_state.wait_with_hidden_states:
+    if global_state.wait_with_hidden_states and not force_hidden_states:
         global_state.local_state.hidden_states = None
     else:
+        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)
         global_state.local_state.hidden_states = hidden_states.cpu().detach()
         
@@ -102,7 +102,7 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     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)
+        get_hidden_states(raw_original_prompt, force_hidden_states=True)
     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