import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch
import pandas as pd

# Load model and tokenizer
model_name = "tscholak/cxmefzzi"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
device = torch.device("cpu")

def generate_sql(natural_language, csv_file):
    # Read uploaded dataset
    if csv_file is not None:
        df = pd.read_csv(csv_file.name)
        table_columns = ", ".join(df.columns)
        context = f"The table has columns: {table_columns}."
    else:
        context = ""

    prompt = f"{context} Convert the question to SQL: {natural_language}"

    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(**inputs, max_length=128)
    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return sql_query

# Gradio UI
iface = gr.Interface(
    fn=generate_sql,
    inputs=[
        gr.Textbox(label="Enter your question (natural language)"),
        gr.File(label="Upload your CSV file (optional)")
    ],
    outputs=gr.Textbox(label="Generated SQL Query"),
    title="🧠 Natural Language to SQL Converter",
    description="Upload a dataset and type your question in natural language to generate an SQL query automatically.",
)

if __name__ == "__main__":
    iface.launch()