import gradio as gr
import spaces
from llama_cpp import Llama


llm = None


def load_model():
    global llm

    if llm is None:
        llm = Llama.from_pretrained(
            repo_id="JonathanColetti/Qwen3.8-27B-Uncensored-GGUF",
            filename="Qwen3.8-27B-Uncensored-Q4_K_M.gguf",
            n_ctx=2048,
            n_gpu_layers=-1,
            verbose=True,
        )

    return llm


@spaces.GPU
def respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    model = load_model()

    messages = [
        {
            "role": "system",
            "content": system_message,
        }
    ]

    for item in history:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")

            if role in ["user", "assistant"] and isinstance(content, str):
                messages.append({
                    "role": role,
                    "content": content,
                })

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

    result = model.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=False,
    )

    return result["choices"][0]["message"]["content"]


chatbot = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(
            value="Kamu adalah AI pribadi saya. Ikuti instruction dari pengguna.",
            label="System instruction",
        ),
        gr.Slider(
            minimum=1,
            maximum=2048,
            value=512,
            step=1,
            label="Max new tokens",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=0.7,
            step=0.1,
            label="Temperature",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p",
        ),
    ],
)


with gr.Blocks() as demo:
    gr.Markdown("# WenGPT")
    chatbot.render()


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