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

# Path to model on the Hugging Face Hub
MODEL_NAME = "Yuk050/gemma-3-1b-text-to-sql-model"

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,  # safer default for CPU Spaces
    device_map="auto",
)

def generate_sql(schema: str, natural_language_query: str) -> str:
    user_message = f"Given the following database schema:\n\n{schema}\n\nGenerate the SQL query for: {natural_language_query}"
    
    messages = [{"role": "user", "content": user_message}]
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    outputs = model.generate(
        input_ids,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
        eos_token_id=tokenizer.eos_token_id
    )

    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
    return response.strip()

iface = gr.Interface(
    fn=generate_sql,
    inputs=[
        gr.Textbox(lines=10, label="Database Schema", placeholder="e.g., CREATE TABLE Employees (...)"),
        gr.Textbox(lines=2, label="Natural Language Query", placeholder="e.g., Select all employees with salary > 50000")
    ],
    outputs=gr.Textbox(lines=5, label="Generated SQL Query"),
    title="Text-to-SQL with Gemma 3 1B",
    description="Enter a database schema and natural language question. The model will generate the SQL."
)

iface.launch()