Update app.py

app.py

@@ -106,6 +106,8 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     if global_state.wait_with_hidden_states and global_state.local_state.hidden_states is None:
         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)
+    hidden_means = torch.tensor(global_state.local_state.hidden_states.mean(dim=1)).to(model.device).to(model.dtype)
+    hidden_norms = hidden_means.norm(dim=-1)
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
     # generation parameters
@@ -131,7 +133,7 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     generation_texts = tokenizer.batch_decode(generated)
 
     # create GUI output
-    important_idxs = 1 + interpreted_vectors.diff(dim=0).norm(dim=-1).topk(k=int(np.ceil(0.2 * len(generation_texts)))).indices.cpu().numpy()
+    important_idxs = 1 + ((interpreted_vectors - hidden_means) / hidden_norms).diff(dim=0).norm(dim=-1).topk(k=int(np.ceil(0.2 * len(generation_texts)))).indices.cpu().numpy()
     print(f'{important_idxs=}')
     progress_dummy_output = ''
     elem_classes = [['bubble', 'even_bubble' if i % 2 == 0 else 'odd_bubble'] +