# app.py — Hugging Face Space (ZeroGPU) Gradio demo
# Fine-tuned Qwen2.5-Coder-7B (LoRA) for Text-to-SQL on the Spider benchmark.

import os
import gradio as gr
import spaces                      # ZeroGPU — provides a free GPU during decorated calls
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
ADAPTER    = "Gansaw98/qwen2.5-coder-7b-text2sql-spider"
HF_TOKEN   = os.environ.get("HF_TOKEN")    # only needed if the adapter repo is private

# Exact system prompt the model was fine-tuned on — do not change.
SYSTEM_PROMPT = (
    "You are an expert SQL generator. "
    "Given a database schema and a natural language question, "
    "write the correct SQL query. "
    "Output only the SQL query with no explanation or markdown."
)

# --- Load once at startup (ZeroGPU provides a GPU for initialization) -----------
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True, token=HF_TOKEN)
base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True, token=HF_TOKEN
).to("cuda")
model = PeftModel.from_pretrained(base, ADAPTER, token=HF_TOKEN)
model.eval()
PAD_ID = tokenizer.pad_token_id or tokenizer.eos_token_id


@spaces.GPU(duration=60)
def generate_sql(schema: str, question: str) -> str:
    schema   = (schema or "").strip()
    question = (question or "").strip()
    if not schema or not question:
        return "-- Please provide both a database schema and a question."

    user = f"### Database Schema:\n{schema}\n\n### Question:\n{question}"
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": user},
    ]
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    ).to(model.device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,        # greedy, deterministic
            num_beams=1,
            pad_token_id=PAD_ID,
        )
    gen = out[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(gen, skip_special_tokens=True).strip()


# --- UI -------------------------------------------------------------------------
EXAMPLE_SCHEMA = """CREATE TABLE singer (
  Singer_ID REAL PRIMARY KEY,
  Name TEXT,
  Country TEXT,
  Song_Name TEXT,
  Song_release_year TEXT,
  Age REAL,
  Is_male TEXT
)

CREATE TABLE concert (
  concert_ID REAL PRIMARY KEY,
  concert_Name TEXT,
  Theme TEXT,
  Stadium_ID TEXT,
  Year TEXT
)"""

with gr.Blocks(title="Text-to-SQL — Fine-tuned Qwen2.5-Coder-7B") as demo:
    gr.Markdown(
        "# 🗃️ Text-to-SQL Demo\n"
        "Fine-tuned **Qwen2.5-Coder-7B** (LoRA / QLoRA) on the **Spider** benchmark — "
        "**77.7% execution accuracy**, beating a zero-shot 70B model by 24.5%.\n\n"
        "Paste a database schema and ask a question in plain English; the model returns SQL."
    )
    with gr.Row():
        with gr.Column():
            schema_in = gr.Textbox(label="Database Schema (CREATE TABLE ...)", lines=14, value=EXAMPLE_SCHEMA)
            question_in = gr.Textbox(label="Question (English)", lines=2,
                                     value="How many singers are there from each country?")
            btn = gr.Button("Generate SQL", variant="primary")
        with gr.Column():
            sql_out = gr.Code(label="Generated SQL", language="sql")

    btn.click(generate_sql, inputs=[schema_in, question_in], outputs=sql_out)

    gr.Examples(
        examples=[
            [EXAMPLE_SCHEMA, "How many singers are there from each country?"],
            [EXAMPLE_SCHEMA, "What are the names of singers older than 40, ordered by age descending?"],
            [EXAMPLE_SCHEMA, "Show the theme and year of every concert."],
        ],
        inputs=[schema_in, question_in],
    )

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