Update app.py

app.py

@@ -96,8 +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!
-        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())
+        hidden_scores = F.normalize(hidden_states, dim=-1).diff(dim=0).norm(dim=-1).cpu() # num_layers x num_tokens
+        important_tokens = (1 + np.unravel_index(hidden_scores.flatten().topk(k=5).indices.numpy(), hidden_scores.shape)[1])
         print(f'{important_tokens=}\t\t{hidden_states.shape=}')
         global_state.local_state.hidden_states = hidden_states.cpu().detach()