Update app.py

app.py

@@ -1,25 +1,39 @@
-import gradio as gr
 import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration
+import gradio as gr
+
+# Load model from local folder
+model_dir = "./"
+
+tokenizer = T5Tokenizer.from_pretrained(model_dir)
+model = T5ForConditionalGeneration.from_pretrained(model_dir)
 
-# Load model and tokenizer from local files
-model_path = "./"
-tokenizer = T5Tokenizer.from_pretrained(model_path)
-model = T5ForConditionalGeneration.from_pretrained(model_path)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+model.to(device)
 model.eval()
 
-def generate_sql(user_input):
-    inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True).to(device)
+# 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)
 
-gr.Interface(
-    fn=generate_sql,
-    inputs=gr.Textbox(lines=15, placeholder="Paste your full SQL prompt here..."),
-    outputs=gr.Textbox(label="Generated SQL Query"),
-    title="HISAB AI – Text2SQL Generator",
-    description="Paste your schema + user query + instructions. Model returns a PostgreSQL query."
-).launch()
+# 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 = 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")
+
+    output = gr.Textbox(label="Generated SQL Query")
+
+    submit = gr.Button("Generate SQL")
+    submit.click(fn=generate_sql, inputs=[schema, instructions, user_query], outputs=output)
+
+demo.launch()