Update app.py

app.py

@@ -30,7 +30,7 @@ fine_tuned_model_id = "bstraehle/Meta-Llama-3-8B-Instruct"
 dataset = "gretelai/synthetic_text_to_sql"
 
 def prompt_model(model_id, system_prompt, user_prompt):
-    pipe = pipeline("text-generation", model=model_id)
+    pipe = pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto")
 
     messages = [
       {"role": "system", "content": system_prompt},
@@ -38,7 +38,7 @@ def prompt_model(model_id, system_prompt, user_prompt):
       {"role": "assistant", "content": ""}
     ]
 
-    output = pipe(messages, model_kwargs={"torch_dtype": torch.bfloat16}, device="cuda")
+    output = pipe(messages)
     
     return output[0]["generated_text"][-1]["content"]
 
@@ -87,7 +87,7 @@ def process(action, system_prompt, user_prompt, schema, base_model_id, fine_tune
     return result
 
 demo = gr.Interface(fn=process, 
-                    inputs=[gr.Radio([action_1, action_2, action_3], label = "Action", value = action_2),
+                    inputs=[gr.Radio([action_1, action_2, action_3], label = "Action", value = action_1),
                             gr.Textbox(label = "System Prompt", value = system_prompt, lines = 2),
                             gr.Textbox(label = "User Prompt", value = user_prompt, lines = 2),
                             gr.Textbox(label = "Schema", value = schema, lines = 2),