NLP-SQL / app.py
Omkar1872's picture
Update app.py
18689a1 verified
raw
history blame
1.71 kB
import gradio as gr
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import sqlite3
# Load model
model_name = "mrm8488/t5-base-finetuned-wikiSQL"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False) # <-- use slow tokenizer
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
def nl_to_sql(question, file):
try:
df = pd.read_csv(file.name)
except Exception as e:
return f"Error reading CSV: {e}", pd.DataFrame()
# Create SQLite DB
conn = sqlite3.connect(":memory:")
df.to_sql("data_table", conn, index=False, if_exists="replace")
# Schema description
schema = ", ".join(df.columns)
text = f"translate English to SQL: {question} | table columns: {schema}"
inputs = tokenizer(text, return_tensors="pt")
outputs = model.generate(**inputs, max_length=256)
sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Execute SQL query
try:
result = pd.read_sql_query(sql_query, conn)
except Exception as e:
result = pd.DataFrame({"Error": [str(e)]})
conn.close()
return sql_query, result.head()
iface = gr.Interface(
fn=nl_to_sql,
inputs=[
gr.Textbox(label="Ask your question (Natural Language)", placeholder="e.g., Show customers older than 30"),
gr.File(label="Upload your CSV file")
],
outputs=[
gr.Textbox(label="Generated SQL Query"),
gr.Dataframe(label="Result Preview")
],
title="🧠 Natural Language to SQL Generator",
description="Upload a CSV and ask questions in plain English. Generates SQL and shows results instantly."
)
if __name__ == "__main__":
iface.launch()