Update app.py

app.py

@@ -99,7 +99,7 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''
     return ([progress_dummy_output] + 
-            [gr.Textbox(text.replace('\n', ' '), visible=True, container=False) for text in generation_texts]
+            [gr.Textbox(text.replace('\n', ' '), visible=True, container=False, label=f'Layer {i}') for text in generation_texts]
            )
 
 
@@ -182,6 +182,9 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
     with gr.Row():
         with gr.Column(scale=5):
             gr.Markdown('# 😎 Self-Interpreting Models')
+
+            gr.Markdown('<b style="color: #8B0000;">Model outputs are not filtered and might include undesired language!</b>')
+            
             # gr.Markdown(
             #     '**👾 This space is a simple introduction to the emerging trend of models interpreting their OWN hidden states in free form natural language!!👾**',
             #     # elem_classes=['explanation_accordion']
@@ -205,7 +208,6 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             we expect to get back a summary of the information that exists inside the hidden state from different layers and different runs!! How cool is that! 😯😯😯
             ''', line_breaks=True)
 
-            gr.Markdown('<b style="color: red;">Model outputs are not filtered and might include undesired language!</b>')
         # with gr.Column(scale=1):    
         #     gr.Markdown('<span style="font-size:180px;">🤔</span>')
 
@@ -238,7 +240,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
     use_gpu = False # gr.Checkbox(value=False, label='Use GPU')
     progress_dummy = gr.Markdown('', elem_id='progress_dummy')
 
-    interpretation_bubbles = [gr.Textbox('', label=f'Layer {i}', container=False, visible=False, elem_classes=['bubble', 
+    interpretation_bubbles = [gr.Textbox('', container=False, visible=False, elem_classes=['bubble', 
                                                                                            'even_bubble' if i % 2 == 0 else 'odd_bubble'])
                              for i in range(model.config.num_hidden_layers)]