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

MODEL_ID = os.getenv("MODEL_ID", "GnLOLot/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking")

model = None
tokenizer = None

@spaces.GPU
def chat_fn(message, history):
    global model, tokenizer
    if model is None:
        print("Loading model...", flush=True)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            trust_remote_code=True
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        print("Model loaded", flush=True)
    messages = []
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id
        )

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

with gr.Blocks(title="MiniCPM5-1B Chat") as demo:
    gr.Markdown(f"# MiniCPM5-1B Chat\n**Model:** `{MODEL_ID}`\n\nPowered by ZeroGPU (free GPU)")
    gr.ChatInterface(
        fn=chat_fn,
        title=None,
        description="First request loads the model (~30s), subsequent calls are faster."
    )

demo.launch()