import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Public model fine-tuned on WikiSQL
model_name = "mrm8488/t5-base-finetuned-wikiSQL"

# Force slow tokenizer to avoid tiktoken issues
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def nl_to_sql(nl_query):
    # Add prefix required for this model
    input_text = "translate English to SQL: " + nl_query
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    outputs = model.generate(inputs, max_length=512)
    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return sql_query

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Natural Language to SQL Generator")
    with gr.Row():
        nl_input = gr.Textbox(label="Enter your query in English")
        sql_output = gr.Textbox(label="Generated SQL")
    btn = gr.Button("Generate SQL")
    btn.click(nl_to_sql, inputs=nl_input, outputs=sql_output)

demo.launch()