Update app.py

app.py

@@ -64,11 +64,11 @@ def prompt_model(model_id, system_prompt, user_prompt, schema):
 #    print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")
 
 def fine_tune_model(base_model_id, dataset):
-    #tokenizer = download_model(base_model_id)
-    prepare_dataset(dataset)
-    train_model(base_model_id)
-    #fine_tuned_model_id = upload_model(base_model_id, tokenizer)
-    return "fine_tuned_model_id"
+    tokenizer = download_model(base_model_id)
+    #prepare_dataset(dataset)
+    #train_model(base_model_id)
+    fine_tuned_model_id = upload_model(base_model_id, tokenizer)
+    return fine_tuned_model_id
 
 def download_model(base_model_id):
     tokenizer = AutoTokenizer.from_pretrained(base_model_id)
@@ -211,7 +211,7 @@ def process(action, base_model_id, dataset, system_prompt, user_prompt, schema):
 
 demo = gr.Interface(fn=process, 
                     inputs=[gr.Radio([action_1, action_2], label = "Action", value = action_1),
-                            gr.Textbox(label = "Base Model ID", value = "bstraehle/code-llama-7b-text-to-sql", lines = 1),
+                            gr.Textbox(label = "Base Model ID", value = base_model_id, lines = 1),
                             gr.Textbox(label = "Dataset", value = dataset, lines = 1),
                             gr.Textbox(label = "System Prompt", value = system_prompt, lines = 2),
                             gr.Textbox(label = "User Prompt", value = user_prompt, lines = 2),