Update app.py

app.py

@@ -13,7 +13,7 @@ MAX_PROMPT_TOKENS = 30
 
 ## info
 dataset_info = [{'name': 'Commonsense', 'hf_repo': 'tau/commonsense_qa', 'text_col': 'question'},
-                {'name': 'Fact Recall', 'hf_repo': 'azhx/counterfact-filtered-gptj6b', 'text_col': 'subject+predicate', 
+                {'name': 'Factual Recall', 'hf_repo': 'azhx/counterfact-filtered-gptj6b', 'text_col': 'subject+predicate', 
                  'filter': lambda x: x['label'] == 1},
                ]
 
@@ -66,7 +66,7 @@ def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
 
 
 def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
-                       temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
+                       temperature, top_k, top_p, repetition_penalty, length_penalty, use_gpu, i, 
                        num_beams=1):
 
     interpreted_vectors = global_state[:, i]
@@ -89,7 +89,8 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
     # generate the interpretations
-    generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
+    generate = spaces.GPU(interpretation_prompt.generate) if use_gpu else interpretation_prompt.generate
+    generated = generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''
     return ([progress_dummy_output] + 
@@ -223,7 +224,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
         for i in range(MAX_PROMPT_TOKENS):
             btn = gr.Button('', visible=False, elem_classes=['token_btn'])
             tokens_container.append(btn)
-    
+    use_gpu = gr.Checkbox(value=True, label='Use GPU')
     progress_dummy = gr.Markdown('', elem_id='progress_dummy')
 
     interpretation_bubbles = [gr.Textbox('', container=False, visible=False, elem_classes=['bubble', 
@@ -255,7 +256,8 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
     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
+                                                     top_k, top_p, repetition_penalty, length_penalty,
+                                                     use_gpu
                                                     ], [progress_dummy, *interpretation_bubbles])
     
     original_prompt_btn.click(get_hidden_states,