Update app.py

app.py

@@ -8,7 +8,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # Model IDs:
 #
-# google/gemma-2-9b-it
 # meta-llama/Meta-Llama-3-8B-Instruct
 
 # Datasets:
@@ -17,9 +16,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 profile = "bstraehle"
 
-action_1 = "Prompt base model"
-action_2 = "Fine-tune base model"
-action_3 = "Prompt fine-tuned model"
+action_1 = "Fine-tune pre-trained model"
+action_2 = "Prompt fine-tuned model"
 
 system_prompt = "You are a text to SQL query translator. Given a question in English, generate a SQL query based on the provided SCHEMA. Do not generate any additional text. SCHEMA: {schema}"
 user_prompt = "What is the total trade value and average price for each trader and stock in the trade_history table?"
@@ -76,12 +74,12 @@ def upload_model(model_id, tokenizer):
 
     return model_repo_name
 
-def process(action, system_prompt, user_prompt, schema, base_model_id, fine_tuned_model_id, dataset):
+def process(action, system_prompt, user_prompt, schema, base_model_id, dataset):
     if action == action_1:
-        result = prompt_model(base_model_id, system_prompt, user_prompt, schema)
-    elif action == action_2:
         result = fine_tune_model(base_model_id)
-    elif action == action_3:
+    elif action == action_2:
+        model_id = base_model_id[base_model_id.rfind('/')+1:]
+        fine_tuned_model_id = f"{profile}/{model_id}"
         result = prompt_model(fine_tuned_model_id, system_prompt, user_prompt, schema)
     
     return result
@@ -92,7 +90,6 @@ demo = gr.Interface(fn=process,
                             gr.Textbox(label = "User Prompt", value = user_prompt, lines = 2),
                             gr.Textbox(label = "Schema", value = schema, lines = 2),
                             gr.Textbox(label = "Base Model ID", value = base_model_id, lines = 1),
-                            gr.Textbox(label = "Fine-Tuned Model ID", value = fine_tuned_model_id, lines = 1),
                             gr.Textbox(label = "Dataset", value = dataset, lines = 1)],
                     outputs=[gr.Textbox(label = "Completion")])
 demo.launch()