Update app.py

app.py

@@ -22,8 +22,10 @@ action_2 = "Prompt fine-tuned model"
 action_3 = "Fine-tune base model"
 
 schema = "CREATE TABLE trade_history (id INT, trader_id INT, stock VARCHAR(255), price DECIMAL(5,2), quantity INT, trade_time TIMESTAMP);"
-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
+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: """
 user_prompt = "What is the total trade value and average price for each trader and stock in the trade_history table?"
 
 base_model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
@@ -77,7 +79,7 @@ def upload_model(model_id, tokenizer):
 
     return model_repo_name
 
-def process(action, system_prompt, user_prompt, base_model_id, fine_tuned_model_id, dataset):
+def process(action, system_prompt, user_prompt, schema, base_model_id, fine_tuned_model_id, dataset):
     if action == action_1:
         result = prompt_model(base_model_id, system_prompt, user_prompt)
     elif action == action_2:
@@ -91,6 +93,7 @@ demo = gr.Interface(fn=process,
                     inputs=[gr.Radio([action_1, action_2, action_3], label = "Action", value = action_2),
                             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),
                             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)],