import gradio as gr
try:
    from transformers import AutoTokenizer, AutoModelForCausalLM
except ImportError:
    import os
    os.system("pip install transformers torch gradio")
    from transformers import AutoTokenizer, AutoModelForCausalLM
    
import torch

# Load model + tokenizer
model_name = "premai-io/prem-1B-SQL"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"  # Uses GPU if available on Spaces
)

def text_to_sql(question, schema=""):
    """
    Convert natural language question into SQL query.
    Schema can be passed as a string (table + column names).
    """
    if schema:
        prompt = f"{schema}\nQuestion: {question}\nSQL:"
    else:
        prompt = f"Question: {question}\nSQL:"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.2,  # Low temp for deterministic SQL
        do_sample=False
    )
    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return sql_query

# Define Gradio interface (API-like, minimal UI)
iface = gr.Interface(
    fn=text_to_sql,
    inputs=[
        gr.Textbox(label="Question"),
        gr.Textbox(label="Schema (optional)", placeholder="table: columns, ...")
    ],
    outputs="text",
    title="Text-to-SQL Converter",
    description="Convert natural language questions into SQL queries using the premai-io/prem-1B-SQL model."
)

# Launch (for Spaces: set share=False, HF will handle the endpoint)
iface.launch()