import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel

BASE_MODEL = "codellama/CodeLlama-7b-Instruct-hf"
ADAPTER_ID = "tamilanda/my-sql-model"

# 1️⃣ Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

# 2️⃣ Force FULL model load first (no meta device)
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
device = 0 if torch.cuda.is_available() else -1

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=dtype,
    low_cpu_mem_usage=False,   # ❗ prevent meta loading
    device_map=None            # ❗ don't auto offload during LoRA merge
)

# 3️⃣ Load LoRA adapter
model = PeftModel.from_pretrained(model, ADAPTER_ID)

# 4️⃣ Move model to correct device
if torch.cuda.is_available():
    model = model.to("cuda")
else:
    model = model.to("cpu")

# 5️⃣ Build pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device
)

# Gradio UI
def generate_sql(prompt):
    if not prompt.strip():
        return "⚠️ Please enter a prompt."
    result = pipe(
        prompt,
        max_new_tokens=150,
        do_sample=True,
        temperature=0.2,
        top_p=0.9,
    )
    return result[0]["generated_text"]

demo = gr.Interface(
    fn=generate_sql,
    inputs=gr.Textbox(label="Enter Natural Language Query"),
    outputs=gr.Textbox(label="Generated SQL"),
    title="🧠 NL → SQL (CodeLlama + LoRA)",
)

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