Update app.py

app.py

@@ -29,11 +29,12 @@ class GlobalState:
     local_state : LocalState = LocalState()
     wait_with_hidden_state : bool = False
     interpretation_prompt_template : str = '{prompt}'
-    original_prompt_template : str = 'User: [X]\n\nAnswer: {prompt}'
+    original_prompt_template : str = 'User: [X]\n\nAssistant: {prompt}'
     layers_format : str = 'model.layers.{k}'
 
 
 suggested_interpretation_prompts = [
+                                    "The meaning of [X] is",
                                     "Sure, here's a bullet list of the key words in your message:",
                                     "Sure, I'll summarize your message:", 
                                     "Sure, here are the words in your message:",
@@ -139,7 +140,7 @@ global_state = GlobalState()
 
 model_name = 'LLAMA2-7B'
 reset_model(model_name, with_extra_components=False)
-original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
+raw_original_prompt = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []
 
 for i in range(MAX_PROMPT_TOKENS):
@@ -185,17 +186,17 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                     dataset = dataset.filter(info['filter'])
                 dataset = dataset.shuffle(buffer_size=2000).take(num_examples)
                 dataset = [[row[info['text_col']]] for row in dataset]
-                gr.Examples(dataset, [original_prompt_raw], cache_examples=False)
+                gr.Examples(dataset, [raw_original_prompt], cache_examples=False)
                 
         with gr.Group():
-            original_prompt_raw.render()
+            raw_original_prompt.render()
             original_prompt_btn = gr.Button('Output Token List', variant='primary')
     
         gr.Markdown('## Choose Your Interpretation Prompt')
         with gr.Group('Interpretation'):
-            interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
+            raw_interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
             interpretation_prompt_examples = gr.Examples([[p] for p in suggested_interpretation_prompts], 
-                                                         [interpretation_prompt], cache_examples=False)
+                                                         [raw_interpretation_prompt], cache_examples=False)
        
         with gr.Accordion(open=False, label='Generation Settings'):
             with gr.Row():
@@ -225,17 +226,17 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
         
     # event listeners    
     for i, btn in enumerate(tokens_container):
-        btn.click(partial(run_interpretation, i=i), [original_prompt_raw, interpretation_prompt, 
+        btn.click(partial(run_interpretation, i=i), [raw_original_prompt, raw_interpretation_prompt, 
                                                      num_tokens, do_sample, temperature, 
                                                      top_k, top_p, repetition_penalty, length_penalty
                                                     ], [progress_dummy, *interpretation_bubbles])
 
     original_prompt_btn.click(get_hidden_states, 
-                              [original_prompt_raw], 
+                              [raw_original_prompt], 
                               [progress_dummy, *tokens_container, *interpretation_bubbles])
-    original_prompt_raw.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)
+    raw_original_prompt.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)
 
-    extra_components = [interpretation_prompt, original_prompt_raw, original_prompt_btn]
+    extra_components = [raw_interpretation_prompt, raw_original_prompt, original_prompt_btn]
     model_chooser.change(reset_model, [model_chooser, *extra_components], 
                          [welcome_model, *interpretation_bubbles, *tokens_container, *extra_components])