Update app.py

app.py

@@ -45,7 +45,7 @@ suggested_interpretation_prompts = ["Before responding, let me repeat the messag
 def initialize_gpu():
     pass
 
-def get_hidden_states(raw_original_prompt):
+def get_hidden_states(progress, raw_original_prompt):
     original_prompt = original_prompt_template.format(prompt=raw_original_prompt)
     model_inputs = tokenizer(original_prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
     tokens = tokenizer.batch_decode(model_inputs.input_ids[0])
@@ -56,7 +56,7 @@ def get_hidden_states(raw_original_prompt):
     return [hidden_states, *token_btns]
 
 
-def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
+def run_interpretation(progress, global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
                        temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
                        num_beams=1):
 
@@ -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 [gr.Text(text, visible=True, container=False) for text in generation_texts]
+    return [gr.TextBox(text, visible=True, container=False) for text in generation_texts]
 
 
 ## main
@@ -137,7 +137,7 @@ css = '''
 
 # '''
 
-
+progress = gr.Progress()
 with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
     global_state = gr.State([])
     with gr.Row():
@@ -147,7 +147,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
                 
                 👾 **This space is a simple introduction to the emerging trend of models interpreting their _own hidden states_ in free form natural language**!! 👾
                 This idea was explored in the paper **Patchscopes** ([Ghandeharioun et al., 2024](https://arxiv.org/abs/2401.06102)) and was later investigated further in **SelfIE** ([Chen et al., 2024](https://arxiv.org/abs/2403.10949)). 
-                An honorary mention of **Speaking Probes** ([Dar, 2023](https://towardsdatascience.com/speaking-probes-self-interpreting-models-7a3dc6cb33d6) -- my own work!! 🥳) which was a less mature but had the same idea in mind. 
+                An honorary mention of **Speaking Probes** ([Dar, 2023](https://towardsdatascience.com/speaking-probes-self-interpreting-models-7a3dc6cb33d6) -- my own work!! 🥳) which was less mature but had the same idea in mind. 
                 We will follow the SelfIE implementation in this space for concreteness. Patchscopes are so general that they encompass many other interpretation techniques too!!! 
                 
                 👾 **The idea is really simple: models are able to understand their own hidden states by nature!** 👾  
@@ -158,7 +158,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             gr.Markdown('<span style="font-size:180px;">🤔</span>')
         
     with gr.Group():
-        original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail', container=True, label='Original Prompt')
+        original_prompt_raw = gr.Textbox(value='Should I eat cake or vegetables?', container=True, label='Original Prompt')
         original_prompt_btn = gr.Button('Compute', variant='primary')
         
     with gr.Accordion(open=False, label='Settings'):
@@ -179,21 +179,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'])
                 tokens_container.append(btn)
-        
-        for i in range(model.config.num_hidden_layers):
-            interpretation_bubbles.append(gr.Text('', container=False, visible=False, elem_classes=['bubble']))
-
+            progress.render()
+            interpretation_bubbles = [gr.TextBox('', container=False, visible=False, elem_classes=['bubble'])
+                                     for i in range(model.config.num_hidden_layers)]
+            
         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
+            btn.click(partial(run_interpretation, i=i), [progress, 
+                                                         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], 
+                              [progress, original_prompt_raw], 
                               [global_state, *tokens_container])    
     demo.launch()