Update app.py

app.py

@@ -131,13 +131,13 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     generation_texts = tokenizer.batch_decode(generated)
 
     # create GUI output
-    print(f'{interpreted_vectors.size()=}\t\t{int(np.ceil(0.2 * len(generation_texts)))=}')
-    important_idxs = 1 + interpreted_vectors.squeeze().diff(dim=0).topk(k=int(np.ceil(0.2 * len(generation_texts))), dim=0).indices.cpu().numpy()
+    important_idxs = 1 + interpreted_vectors.diff(dim=0).topk(k=int(np.ceil(0.2 * len(generation_texts))), dim=0).indices.cpu().numpy()
+    print(f'{important_idxs=}')
     progress_dummy_output = ''
     elem_classes = [['bubble', 'even_bubble' if i % 2 == 0 else 'odd_bubble'] + 
                     ([] if i in important_idxs else ['faded_bubble']) for i in range(len(generation_texts))]
     bubble_outputs = [gr.Textbox(text.replace('\n', ' '), show_label=True, visible=True, 
-                                 container=False, label=f'Layer {i}', elem_classes=elem_classes[i])
+                                 container=True, label=f'Layer {i}', elem_classes=elem_classes[i])
                       for i, text in enumerate(generation_texts)]
     bubble_outputs += [gr.Textbox('', visible=False) for _ in range(MAX_NUM_LAYERS - len(bubble_outputs))]
     return [progress_dummy_output, *bubble_outputs]