Update app.py

app.py

@@ -11,6 +11,7 @@ from interpret import InterpretationPrompt
 
 MAX_PROMPT_TOKENS = 60
 
+
 ## info
 dataset_info = [
                 {'name': 'Commonsense', 'hf_repo': 'tau/commonsense_qa', 'text_col': 'question'},
@@ -56,7 +57,7 @@ suggested_interpretation_prompts = ["Sure, I'll summarize your message:", "Sure,
 def initialize_gpu():
     pass
 
-def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
+def get_hidden_states(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])
@@ -65,17 +66,13 @@ def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
     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 = ''
-    return [progress_dummy_output, hidden_states, *token_btns]
-
-
-@spaces.GPU
-def generate_interpretation_gpu(interpret_prompt, *args, **kwargs):
-    return interpret_prompt.generate(*args, **kwargs)
+    invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_PROMPT_TOKENS)]
+    return [progress_dummy_output, hidden_states, *token_btns, *invisible_bubbles]
 
 
 @spaces.GPU
 def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
-                       temperature, top_k, top_p, repetition_penalty, length_penalty, use_gpu, i, 
+                       temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
                        num_beams=1):
 
     interpreted_vectors = global_state[:, i]
@@ -98,8 +95,8 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
     # generate the interpretations
-    generate = generate_interpretation_gpu if use_gpu else lambda interpretation_prompt, *args, **kwargs: interpretation_prompt.generate(*args, **kwargs)
-    generated = generate(interpretation_prompt, model, {0: interpreted_vectors}, k=3, **generation_kwargs)
+    # generate = generate_interpretation_gpu if use_gpu else lambda interpretation_prompt, *args, **kwargs: interpretation_prompt.generate(*args, **kwargs)
+    generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''
     return ([progress_dummy_output] + 
@@ -187,14 +184,14 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
             
     with gr.Group():
         original_prompt_raw.render()
-        original_prompt_btn = gr.Button('Compute', variant='primary')
+        original_prompt_btn = gr.Button('Output Token List', variant='primary')
 
     tokens_container = []
     with gr.Row():
+        gr.Markdown('### Here go the tokens of the prompt (click on the one to explore)')
         for i in range(MAX_PROMPT_TOKENS):
             btn = gr.Button('', visible=False, elem_classes=['token_btn'])
             tokens_container.append(btn)
-    use_gpu = False # gr.Checkbox(value=False, label='Use GPU')
     progress_dummy = gr.Markdown('', elem_id='progress_dummy')
 
     
@@ -226,12 +223,12 @@ 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, use_gpu=use_gpu), [global_state, interpretation_prompt, 
+        btn.click(partial(run_interpretation, i=i), [global_state, 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], 
-                              [progress_dummy, global_state, *tokens_container])    
+                              [progress_dummy, global_state, *tokens_container, *interpretation_bubbles])    
     demo.launch()