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. Users will ask you a question in English and you will generate a SQL query."
+system_prompt = "You are a text to SQL query translator. A user will ask you a question in English and you will generate a SQL query. Do not generate any additional text."
 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"
@@ -36,8 +36,10 @@ def prompt_model(model_id, system_prompt, user_prompt):
       {"role": "user", "content": user_prompt},
       {"role": "assistant", "content": ""}
     ]
-        
-    return pipe(messages)
+
+    outputs = pipe(messages)
+    
+    return outputs[0]["generated_text"][-1]
 
 def fine_tune_model(model_id):
     tokenizer = download_model(model_id)