Update app.py

app.py

@@ -1,39 +1,53 @@
-import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
 import gradio as gr
+import torch
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 
-# Load model from local folder
-model_dir = "./"
-
-tokenizer = T5Tokenizer.from_pretrained(model_dir)
-model = T5ForConditionalGeneration.from_pretrained(model_dir)
-
+# Load tokenizer and model from local folder
+model_path = "./"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+
+tokenizer = T5Tokenizer.from_pretrained(model_path)
+model = T5ForConditionalGeneration.from_pretrained(model_path).to(device)
 model.eval()
 
-# Inference function
-def generate_sql(schema, instructions, user_query):
-    combined_input = f"{instructions.strip()}\n\n{schema.strip()}\n\nUser Query: \"{user_query.strip()}\"\n\nSQL Query:"
-    inputs = tokenizer(combined_input, padding=True, truncation=True, return_tensors="pt").to(device)
-    
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=512)
-    
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_sql(schema, user_query, date_info=""):
+    prompt = f"""Given the following database schema and requirements, generate a PostgreSQL query:
 
-# UI layout
-with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Text-to-SQL Generator")
-    gr.Markdown("Enter the **schema**, **prompt/instructions**, and a **user query** to get the SQL output.")
+{schema}
 
-    schema = gr.Textbox(label="Database Schema", lines=10, placeholder="CREATE TABLE students (...) ...")
-    instructions = gr.Textbox(label="SQL Instructions / Prompt", lines=15, placeholder="Explain how to generate SQL queries...")
-    user_query = gr.Textbox(label="User Query", placeholder="e.g., Show me students who never attended class")
+User Query: "{user_query}"
 
-    output = gr.Textbox(label="Generated SQL Query")
+IMPORTANT REQUIREMENTS:
+- Always filter by user_id = $1 for security
+- The current year is 2025. You are working in this year!
+- CRITICAL: Use ONLY the dates provided in the input parameters. Do NOT infer or change dates on your own!
+- Single date provided: {date_info if date_info else '[none provided]'}
+- If date range is provided, use DATE(created_at) BETWEEN 'startDate' AND 'endDate'
+- If single date is provided, use DATE(created_at) = 'YYYY-MM-DD'
+- NEVER use hardcoded years like 2024 - always use the provided dates exactly as given
+- Generate ONLY simple SELECT, INSERT, UPDATE, DELETE statements.
+- NO WITH clauses, NO CTEs, NO complex subqueries.
+- Use RETURNING * for INSERT, UPDATE, DELETE when applicable.
 
-    submit = gr.Button("Generate SQL")
-    submit.click(fn=generate_sql, inputs=[schema, instructions, user_query], outputs=output)
+SQL Query:"""
 
-demo.launch()
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True).to(device)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=256)
+    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_sql
+
+# Gradio UI
+iface = gr.Interface(
+    fn=generate_sql,
+    inputs=[
+        gr.Textbox(label="Database Schema", lines=12, placeholder="CREATE TABLE ..."),
+        gr.Textbox(label="User Query", placeholder="How much did I spend on food last week?"),
+        gr.Textbox(label="Date Info (optional)", placeholder="2025-06-12 or 2025-06-01 to 2025-06-07")
+    ],
+    outputs=gr.Textbox(label="Generated SQL Query"),
+    title="HISAB AI - Natural Language to SQL",
+    description="Enter your schema, user query and date (optional). Model will output SQL query."
+)
+
+iface.launch()