Update app.py

app.py

@@ -1,44 +1,38 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load text-to-SQL model
+# Load model — stick with this one for now, later we can upgrade
 generator = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-wikiSQL")
 
-# Function to process full JSON input
 def convert_to_sql(payload):
     question = payload.get("question", "")
     schema = payload.get("schema", "")
     sample_rows = payload.get("sample_rows", [])
 
-    # Turn sample data into a preview string
-    preview_rows = "\n".join([str(row) for row in sample_rows[:5]])
-
-    # Build enhanced prompt
+    # Craft prompt
     prompt = f"""
-You are an AI that writes DuckDB-compatible SQL based on user questions.
-
-Table Schema:
+You are an AI that converts natural language into SQL for a DuckDB database.
+Given a table with the following schema:
 {schema}
 
-Sample Data:
-{preview_rows}
+Here are some sample rows:
+{sample_rows}
 
-Question:
-{question}
+Write a syntactically correct SQL query (DuckDB-compatible) to answer this question: "{question}"
 
-Write only the SQL query — no explanations.
-""".strip()
+Only return the SQL query — no explanation, no markdown.
+    """.strip()
 
-    result = generator(prompt, max_length=256, clean_up_tokenization_spaces=True)[0]["generated_text"]
+    result = generator(prompt, max_length=256)[0]["generated_text"]
     return result.strip()
 
-# Gradio interface for testing (takes JSON input)
+# Define inputs/outputs for interactive mode (not used by FastAPI)
 demo = gr.Interface(
     fn=convert_to_sql,
     inputs=gr.JSON(label="question + schema + sample_rows"),
     outputs="text",
-    title="Smarter Text-to-SQL Generator",
-    description="Send in natural question, schema, and data — get back accurate SQL."
+    title="Text-to-SQL Generator (DuckDB)",
+    description="Send a JSON payload with question, schema, and sample rows"
 )
 
 demo.launch()