import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "melihemin/qwen2.5-0.5b-text2sql-full"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto"
)

def text_to_sql(question):
    prompt = f"""### Question:
{question}

### SQL:
"""
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def load_example():
    return "How many heads of the departments are older than 56?"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Melih Emin Kılıçoğlu Text-to-SQL Demo  
        This application converts **natural language questions** into **SQL queries**  
        using a fine-tuned **Qwen2.5-0.5B** language model.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_box = gr.Textbox(
                label="📝 Natural Language Question",
                placeholder="Enter your question here...",
                lines=8
            )

            generate_btn = gr.Button("🚀 Generate SQL", variant="primary")

        with gr.Column(scale=1):
            output_box = gr.Textbox(
                label="🧾 Generated SQL Query",
                lines=10,
                interactive=False
            )

    generate_btn.click(
        fn=text_to_sql,
        inputs=input_box,
        outputs=output_box
    )

    gr.Markdown("---")

    gr.Markdown(
        """
        ## 📌 Example Usage  
        Click the button below to load a sample question and test the model.
        """
    )

    example_btn = gr.Button("📎 Load Example Question")
    example_btn.click(
        fn=load_example,
        inputs=None,
        outputs=input_box
    )

    gr.Markdown(
        """
        **Example Question:**  
        *How many heads of the departments are older than 56?*

        OR

        *Tüm öğrencileri listele

        **Expected SQL Output:**  
        ```sql
        SELECT count(*) FROM head WHERE age > 56;
        ```

        OR

        ```sql
        SELECT * from students;
        ```
        """
    )

demo.launch()