File size: 2,159 Bytes
dec6df4 a0ff27a dec6df4 a0ff27a 82a8738 a0ff27a 82a8738 a0ff27a 82a8738 a0ff27a dec6df4 a0ff27a dec6df4 a0ff27a dec6df4 a0ff27a dec6df4 a0ff27a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import gradio as gr
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
# Load tokenizer and model from local folder
model_path = "./"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = T5Tokenizer.from_pretrained(model_path)
model = T5ForConditionalGeneration.from_pretrained(model_path).to(device)
model.eval()
def generate_sql(schema, user_query, date_info=""):
prompt = f"""Given the following database schema and requirements, generate a PostgreSQL query:
{schema}
User Query: "{user_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.
SQL Query:"""
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()
|