Update app.py

app.py

@@ -96,9 +96,9 @@ 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 = set(1 + F.normalize(hidden_states, dim=-1)
-                               .diff(dim=0).norm(dim=-1)
-                               .topk(k=7).indices[:, 0].cpu().numpy())
+        important_tokens = (1 + 
+                            F.normalize(hidden_states, dim=-1).diff(dim=0)
+                            .norm(dim=-1).topk(k=7).indices[:, 0].cpu().numpy())
         print(f'{important_tokens=}')
         global_state.local_state.hidden_states = hidden_states.cpu().detach()