Update app.py

app.py

@@ -82,7 +82,7 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     # generate the interpretations
     generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
-    return generation_texts
+    return [gr.Text(text, visible=True) for text in generation_texts]
 
 
 ## main
@@ -103,7 +103,14 @@ tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, token=os.environ['hf_t
 
 # demo
 json_output = gr.JSON()
-css = ''
+css = '''
+.bubble {
+  border: 2px solid #000;
+  border-radius: 10px;
+  padding: 10px;
+}
+'''
+
 
 # '''
 # .token_btn{
@@ -165,14 +172,21 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
 
     with gr.Group('Output'):
         tokens_container = []
+        interpretation_bubbles = []
         with gr.Row():
             for i in range(MAX_PROMPT_TOKENS):
                 btn = gr.Button('', visible=False, elem_classes=['token_btn'])
-                btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, num_tokens, do_sample, temperature, 
-                                                           top_k, top_p, repetition_penalty, length_penalty
-                                                          ], [json_output])
                 tokens_container.append(btn)
-        json_output.render()
         
-    original_prompt_btn.click(get_hidden_states, [original_prompt_raw], [global_state, *tokens_container])    
+        for i in range(model.config.num_hidden_layers):
+            interpretation_bubbles.append(gr.Text('', visible=False, elrm_classes=['bubble']))
+
+        for i, btn in enumerate(tokens_container):
+            btn.click(partial(run_interpretation, i=i), [global_state, interpretation_prompt, num_tokens, do_sample, temperature, 
+                                                               top_k, top_p, repetition_penalty, length_penalty
+                                                              ], [*interpretation_bubbles])
+        
+    original_prompt_btn.click(get_hidden_states, 
+                              [original_prompt_raw], 
+                              [global_state, *tokens_container])    
     demo.launch()