nl2sql-api / app.py
Chaitanya182004's picture
Update app.py
ca9941d verified
Raw
History Blame Contribute Delete
2.24 kB
import os
# Disable Gradio SSR (better for Render)
os.environ["GRADIO_SSR_MODE"] = "False"
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
MODEL_NAME = "gaussalgo/T5-LM-Large-text2sql-spider"
tokenizer = None
model = None
device = "cuda" if torch.cuda.is_available() else "cpu"
def load_model():
global tokenizer, model
if model is None:
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME
)
model = AutoModelForSeq2SeqLM.from_pretrained(
MODEL_NAME
)
model.to(device)
model.eval()
print(f"Model ready on {device}")
def generate_sql(question, context):
try:
load_model()
input_text = f"{question} | {context}"
inputs = tokenizer(
input_text,
return_tensors="pt",
max_length=512,
truncation=True
).to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=128,
num_beams=4,
early_stopping=True
)
sql = tokenizer.decode(
outputs[0],
skip_special_tokens=True
)
return sql
except Exception as e:
return f"Error: {str(e)}"
with gr.Blocks() as demo:
gr.Markdown(
"# NL2SQL API\nGenerate SQL from Natural Language"
)
with gr.Row():
question = gr.Textbox(
label="Question",
placeholder="Example: Find all users"
)
context = gr.Textbox(
label="Database Schema",
placeholder="Example: users(id,name,email)"
)
output = gr.Textbox(
label="Generated SQL",
lines=5
)
btn = gr.Button(
"Generate SQL"
)
btn.click(
fn=generate_sql,
inputs=[
question,
context
],
outputs=output
)
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=int(
os.environ.get("PORT",7860)
),
share=False
)