Update app.py

app.py

@@ -96,9 +96,8 @@ def get_hidden_states(raw_original_prompt, force_hidden_states=False):
         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)
         # TODO: document this!
-        important_tokens = (1 + 
-                            F.normalize(hidden_states, dim=-1).diff(dim=0)
-                            .norm(dim=-1).topk(k=4).indices[:, 0].cpu().numpy())
+        hidden_scores = F.normalize(hidden_states, dim=-1).diff(dim=0).norm(dim=-1)
+        important_tokens = (1 + hidden_scores.topk(k=4).indices[:, 1].cpu().numpy())
         print(f'{important_tokens=}\t\t{hidden_states.shape=}')
         global_state.local_state.hidden_states = hidden_states.cpu().detach()