Update app.py

app.py

@@ -37,8 +37,7 @@ def prompt_model(model_id, system_prompt, user_prompt):
       {"role": "assistant", "content": ""}
     ]
 
-    output = pipe(messages)
-    print(output)
+    output = pipe(messages, device="cuda")
     
     return output[0]["generated_text"][-1]["content"]
 
@@ -46,7 +45,7 @@ def fine_tune_model(model_id):
     tokenizer = download_model(model_id)
     model_repo_name = upload_model(model_id, tokenizer)
 
-    return ""
+    return model_repo_name
         
 def download_model(model_id):
     tokenizer = AutoTokenizer.from_pretrained(model_id)