Update app.py

app.py

@@ -8,22 +8,22 @@ ACTION_1 = "Prompt base model"
 ACTION_2 = "Fine-tune base model"
 ACTION_3 = "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?"
-SQL_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 SQL_CONTEXT. Do not generate any additional text. SQL_CONTEXT: {sql_context}"
+USER_PROMPT = "How many new users joined from countries with stricter data privacy laws than the United States in the past month?"
+SQL_CONTEXT = "CREATE TABLE users (user_id INT, country VARCHAR(50), joined_date DATE); CREATE TABLE data_privacy_laws (country VARCHAR(50), privacy_level INT); INSERT INTO users (user_id, country, joined_date) VALUES (1, 'USA', '2023-02-15'), (2, 'Germany', '2023-02-27'); INSERT INTO data_privacy_laws (country, privacy_level) VALUES ('USA', 5), ('Germany', 8);"
 
 BASE_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 FT_MODEL_NAME = "bstraehle/Meta-Llama-3.1-8B-Instruct-text-to-sql"
 DATASET_NAME = "gretelai/synthetic_text_to_sql"
 
-def process(action, base_model_name, ft_model_name, dataset_name, system_prompt, user_prompt, sql_schema):
+def process(action, base_model_name, ft_model_name, dataset_name, system_prompt, user_prompt, sql_context):
     #raise gr.Error("Please clone and bring your own credentials.")
     if action == ACTION_1:
-        result = prompt_model(base_model_name, system_prompt, user_prompt, sql_schema)
+        result = prompt_model(base_model_name, system_prompt, user_prompt, sql_context)
     elif action == ACTION_2:
         result = fine_tune_model(base_model_name, dataset_name)
     elif action == ACTION_3:
-        result = prompt_model(ft_model_name, system_prompt, user_prompt, sql_schema)
+        result = prompt_model(ft_model_name, system_prompt, user_prompt, sql_context)
     return result
 
 def fine_tune_model(base_model_name, dataset_name):
@@ -108,14 +108,14 @@ def fine_tune_model(base_model_name, dataset_name):
     # Train model
     trainer.train()
     
-def prompt_model(model_name, system_prompt, user_prompt, sql_schema):
+def prompt_model(model_name, system_prompt, user_prompt, sql_context):
     pipe = pipeline("text-generation", 
                     model=model_name,
                     device_map="auto",
                     max_new_tokens=1000)
     
     messages = [
-      {"role": "system", "content": system_prompt.format(schema=sql_schema)},
+      {"role": "system", "content": system_prompt.format(sql_context=sql_context)},
       {"role": "user", "content": user_prompt},
       {"role": "assistant", "content": ""}
     ]
@@ -144,6 +144,6 @@ demo = gr.Interface(fn=process,
                             gr.Textbox(label = "Dataset Name", value = DATASET_NAME, lines = 1),
                             gr.Textbox(label = "System Prompt", value = SYSTEM_PROMPT, lines = 2),
                             gr.Textbox(label = "User Prompt", value = USER_PROMPT, lines = 2),
-                            gr.Textbox(label = "SQL Schema", value = SQL_SCHEMA, lines = 2)],
+                            gr.Textbox(label = "SQL Context", value = SQL_CONTEXT, lines = 2)],
                     outputs=[gr.Textbox(label = "Prompt Completion", value = os.environ["OUTPUT"])])
 demo.launch()