Update app.py

app.py

@@ -70,7 +70,7 @@ def get_hidden_states(raw_original_prompt):
     token_btns = ([gr.Button(token, visible=True) for token in tokens] 
                   + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
     progress_dummy_output = ''
-    invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_PROMPT_TOKENS)]
+    invisible_bubbles = [gr.Textbox('', visible=False) for i in range(len(interpretation_bubbles))]
     return [progress_dummy_output, hidden_states, *token_btns, *invisible_bubbles]
 
 
@@ -126,8 +126,11 @@ model = AutoModelClass.from_pretrained(model_path, **model_args).cuda()
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
 
 # demo
-
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
+tokens_container = []
+        for i in range(MAX_PROMPT_TOKENS):
+            btn = gr.Button('', visible=False, elem_classes=['token_btn'])
+            tokens_container.append(btn)
 
 with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     global_state = gr.State([])
@@ -155,9 +158,9 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
             gr.Markdown(
             '''
             **👾 The idea is really simple: models are able to understand their own hidden states by nature! 👾**
-            According to the residual stream view ([nostalgebraist, 2020](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)), internal representations from different layers are transferable between layers. 
-            So we can inject an representation from (roughly) any layer to any layer! If I give a model a prompt of the form ``User: [X] Assistant: Sure'll I'll repeat your message`` and replace the internal representation of ``[X]`` *during computation* with the hidden state we want to understand, 
-            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! 😯😯😯
+            In line with the residual stream view ([nostalgebraist, 2020](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)), internal representations from different layers are transferable between layers. 
+            So we can inject an representation from (roughly) any layer into any layer! If we give a model a prompt of the form ``User: [X] Assistant: Sure'll I'll repeat your message`` and replace the internal representation of ``[X]`` *during computation* with the hidden state we want to understand, 
+            we expect to get back a summary of the information that exists inside the hidden state, despite being from a different layer and a different run!! How cool is that! 😯😯😯
             ''', line_breaks=True)
 
         # with gr.Column(scale=1):    
@@ -183,21 +186,19 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
             if 'filter' in info:
                 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)
+            dataset = [[row[info['text_col']] + [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)]] for row in dataset]
+            gr.Examples(dataset, [original_prompt_raw, *tokens_container], cache_examples=False)
             
     with gr.Group():
         original_prompt_raw.render()
         original_prompt_btn = gr.Button('Output Token List', variant='primary')
 
-    tokens_container = []
     gr.Markdown('### Here go the tokens of the prompt (click on the one to explore)')
+    
     with gr.Row():
-        for i in range(MAX_PROMPT_TOKENS):
-            btn = gr.Button('', visible=False, elem_classes=['token_btn'])
-            tokens_container.append(btn)
-    progress_dummy = gr.Markdown('', elem_id='progress_dummy')
-
+        for btn in tokens_container:
+            btn.render()
+    
     with gr.Accordion(open=False, label='Generation Settings'):
         with gr.Row():
             num_tokens = gr.Slider(1, 100, step=1, value=20, label='Max. # of Tokens')
@@ -210,8 +211,8 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
                 temperature = gr.Slider(0., 5., value=0.6, label='Temperature')
                 top_k = gr.Slider(1, 1000, value=50, step=1, label='top k')
                 top_p = gr.Slider(0., 1., value=0.95, label='top p')
-
-
+    
+    progress_dummy = gr.Markdown('', elem_id='progress_dummy')
     
     interpretation_bubbles = [gr.Textbox('', container=False, visible=False, elem_classes=['bubble', 
                                                                                            'even_bubble' if i % 2 == 0 else 'odd_bubble'])