Update app.py

app.py

@@ -104,11 +104,15 @@ def get_hidden_states(raw_original_prompt, force_hidden_states=False):
 
 @spaces.GPU
 def run_interpretation(raw_original_prompt, 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):
     model = global_state.model
     tokenizer = global_state.tokenizer
     print(f'run {model}')
+    if use_gpu:
+        model = model.cuda()
+    else:
+        model = model.cpu()
     if global_state.wait_with_hidden_states and global_state.local_state.hidden_states is None:
         get_hidden_states(raw_original_prompt, force_hidden_states=True)
     interpreted_vectors = torch.tensor(global_state.local_state.hidden_states[:, i]).to(model.device).to(model.dtype)
@@ -251,6 +255,7 @@ with gr.Blocks(theme=gr.themes.Glass(), css='styles.css') as demo:
         with gr.Row():
             for btn in tokens_container:
                 btn.render()
+        use_gpu = gr.Radio('Use GPU', value=True)
         
         progress_dummy = gr.Markdown('', elem_id='progress_dummy')
         interpretation_bubbles = [gr.Textbox('', container=False, visible=False) for i in range(MAX_NUM_LAYERS)]
@@ -259,7 +264,8 @@ with gr.Blocks(theme=gr.themes.Glass(), css='styles.css') as demo:
     for i, btn in enumerate(tokens_container):
         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
+                                                     top_k, top_p, repetition_penalty, length_penalty,
+                                                     use_gpu
                                                     ], [progress_dummy, *interpretation_bubbles])
 
     original_prompt_btn.click(get_hidden_states,