import spaces
import gradio as gr
import torch

from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "Qwen/Qwen2.5-Coder-1.5B-Instruct"

print("Loading original Qwen model...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    use_safetensors=True,
    low_cpu_mem_usage=True
).to("cuda")

model.eval()

print("Qwen model loaded successfully!")


@spaces.GPU(duration=60)
def chat(message, history):

    messages = [
        {
            "role": "system",
            "content": "You are an expert AI coding assistant."
        }
    ]

    for item in history:

        # Newer Gradio message format
        if isinstance(item, dict):
            if item.get("role") in ["user", "assistant"]:
                messages.append({
                    "role": item["role"],
                    "content": item["content"]
                })

        # Older Gradio tuple format
        elif isinstance(item, (tuple, list)) and len(item) == 2:
            user_msg, assistant_msg = item

            if user_msg:
                messages.append({
                    "role": "user",
                    "content": user_msg
                })

            if assistant_msg:
                messages.append({
                    "role": "assistant",
                    "content": assistant_msg
                })

    messages.append({
        "role": "user",
        "content": message
    })

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=4096
    ).to("cuda")

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.6,
            do_sample=True,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )

    generated = output[0][inputs["input_ids"].shape[1]:]

    return tokenizer.decode(
        generated,
        skip_special_tokens=True
    )


demo = gr.ChatInterface(
    fn=chat,
    title="Qwen2.5 Coder",
    description="Original Qwen model powered by Hugging Face ZeroGPU.",
    examples=[
        "Write a Python binary search function.",
        "Build a FastAPI application.",
        "Explain Python decorators.",
        "Debug this Python code."
    ]
)

demo.launch()