Update app.py

app.py

@@ -95,9 +95,16 @@ def get_hidden_states(raw_original_prompt, force_hidden_states=False):
     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)
+        # TODO: document this!
+        important_tokens = set(1 + F.normalize(hidden_states, dim=-1)
+                               .diff(dim=0).norm(dim=-1)
+                               .topk(k=10, dim=0).values.topk(k=5).indices.cpu().numpy())
         global_state.local_state.hidden_states = hidden_states.cpu().detach()
         
-    token_btns = ([gr.Button(token, visible=True) for token in tokens] 
+    token_btns = ([gr.Button(token, visible=True, 
+                             elem_classes=['token_btn'] + (['important_token'] if i in important_tokens else [])
+                            ) 
+                   for i, token in enumerate(tokens)] 
                   + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
     progress_dummy_output = ''
     invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_NUM_LAYERS)]
@@ -182,7 +189,7 @@ raw_original_prompt = gr.Textbox(value='How to make a Molotov cocktail?', contai
 tokens_container = []
 
 for i in range(MAX_PROMPT_TOKENS):
-    btn = gr.Button('', visible=False, elem_classes=['token_btn'])
+    btn = gr.Button('', visible=False)
     tokens_container.append(btn)
 
 with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo: