Update app.py

app.py

@@ -52,7 +52,7 @@ suggested_interpretation_prompts = [
 def initialize_gpu():
     pass
 
-def reset_model(model_name, *extra_components, reset_sentence_transformer=False, with_extra_components=True): 
+def reset_model(model_name, load_on_gpu, *extra_components, reset_sentence_transformer=False, with_extra_components=True): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -72,7 +72,7 @@ def reset_model(model_name, *extra_components, reset_sentence_transformer=False,
     if reset_sentence_transformer:
         global_state.sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')
         gc.collect()
-    if not dont_cuda:
+    if  load_on_gpu and not dont_cuda:
         global_state.model.to('cuda')
     global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
     gc.collect()
@@ -175,7 +175,7 @@ torch.set_grad_enabled(False)
 global_state = GlobalState()
 
 model_name = 'LLAMA2-7B'
-reset_model(model_name, with_extra_components=False, reset_sentence_transformer=True)
+reset_model(model_name, load_on_gpu=True, with_extra_components=False, reset_sentence_transformer=True)
 raw_original_prompt = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []
 
@@ -211,6 +211,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
 
     with gr.Group():
         model_chooser = gr.Radio(label='Choose Your Model', choices=list(model_info.keys()), value=model_name)
+        load_on_gpu = gr.Checkbox(label='Load on GPU', value=True)
         welcome_model = gr.Markdown(welcome_message.format(model_name=model_name))
     with gr.Blocks() as demo_main:
         gr.Markdown('## The Prompt to Analyze')        
@@ -274,7 +275,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     raw_original_prompt.change(lambda: [gr.Button(visible=False) for _ in range(MAX_PROMPT_TOKENS)], [], tokens_container)
 
     extra_components = [raw_interpretation_prompt, raw_original_prompt, original_prompt_btn]
-    model_chooser.change(reset_model, [model_chooser, *extra_components], 
+    model_chooser.change(reset_model, [model_chooser, load_on_gpu, *extra_components], 
                          [welcome_model, *interpretation_bubbles, *tokens_container, *extra_components])
     
     demo.launch()