from pathlib import Path

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


BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"
ADAPTER_DIR = Path(__file__).resolve().parent


def load_model():
    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_DIR)
    use_cuda = torch.cuda.is_available()
    dtype = torch.float16 if use_cuda else torch.float32
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
    )
    if use_cuda:
        # ZeroGPU requires the model to be placed on CUDA during app startup.
        base_model = base_model.to("cuda")
    else:
        base_model = base_model.to("cpu")
    model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
    model.eval()
    return model, tokenizer


model, tokenizer = load_model()


def gpu_duration(prompt: str, max_new_tokens: int, temperature: float) -> int:
    # Allow time for long requests while keeping typical calls in the fast queue.
    return min(120, max(30, 15 + int(max_new_tokens) // 4))


@spaces.GPU(duration=gpu_duration)
@torch.inference_mode()
def predict(prompt: str, max_new_tokens: int, temperature: float) -> str:
    if not prompt.strip():
        raise gr.Error("Enter a prompt before requesting an inference.")

    inputs = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    inputs = inputs.to(model.get_input_embeddings().weight.device)
    generation_options = {
        "max_new_tokens": int(max_new_tokens),
        "eos_token_id": tokenizer.eos_token_id,
        "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
        "do_sample": temperature > 0,
    }
    if temperature > 0:
        generation_options["temperature"] = temperature

    output = model.generate(**inputs, **generation_options)
    generated_tokens = output[0, inputs["input_ids"].shape[1] :]
    return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()


demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Prompt", lines=5, placeholder="Ask for help writing or explaining code..."),
        gr.Slider(minimum=16, maximum=512, value=128, step=16, label="Maximum new tokens"),
        gr.Slider(minimum=0, maximum=1.5, value=0.7, step=0.1, label="Temperature (0 = deterministic)"),
    ],
    outputs=gr.Textbox(label="Model response", lines=10),
    title="GRPO Qwen2.5 Coder",
    description="Qwen2.5-Coder-3B-Instruct fine-tuned with a GRPO LoRA adapter.",
    api_name="predict",
)

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