Update app.py

app.py

@@ -21,7 +21,7 @@ action_1 = "Prompt base model"
 action_2 = "Prompt fine-tuned model"
 action_3 = "Fine-tune base 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: "
+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?"
 schema = "CREATE TABLE trade_history (id INT, trader_id INT, stock VARCHAR(255), price DECIMAL(5,2), quantity INT, trade_time TIMESTAMP);"
 
@@ -29,15 +29,17 @@ base_model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 fine_tuned_model_id = "bstraehle/Meta-Llama-3-8B-Instruct"
 dataset = "gretelai/synthetic_text_to_sql"
 
-def prompt_model(model_id, system_prompt, user_prompt):
+def prompt_model(model_id, system_prompt, user_prompt, schema):
     pipe = pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto")
 
     messages = [
-      {"role": "system", "content": system_prompt},
+      {"role": "system", "content": system_prompt.format(schema=schema)},
       {"role": "user", "content": user_prompt},
       {"role": "assistant", "content": ""}
     ]
 
+    print(messages)
+    
     output = pipe(messages)
     
     return output[0]["generated_text"][-1]["content"]
@@ -78,9 +80,9 @@ def upload_model(model_id, tokenizer):
 
 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)
+        result = prompt_model(base_model_id, system_prompt, user_prompt, schema)
     elif action == action_2:
-        result = prompt_model(fine_tuned_model_id, system_prompt, user_prompt)
+        result = prompt_model(fine_tuned_model_id, system_prompt, user_prompt, schema)
     elif action == action_3:
         result = fine_tune_model(base_model_id)