import gradio as gr
import spaces
from llama_cpp import Llama
import os

from huggingface_hub import hf_hub_download

MODEL_PATH = hf_hub_download(
    repo_id="im-lemon/tinychatbot-1b-v1.0",
    filename="tinychatbot-1b-epoch-1.Q4_K_M.gguf"
)

llm = None


def load_model():
    global llm

    if llm is None:
        print("Loading model...")

        llm = Llama(
            model_path=MODEL_PATH,
            n_ctx=2048,
            n_gpu_layers=-1,
            verbose=True
        )

        print("Model loaded!")

    return llm


@spaces.GPU
def chat(message, history):
    model = load_model()

    prompt = ""

    for user_msg, bot_msg in history:
        prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"

    prompt += f"User: {message}\nAssistant:"

    response = model(
        prompt,
        max_tokens=256,
        temperature=0.7,
        stop=["User:"]
    )

    return response["choices"][0]["text"].strip()


demo = gr.ChatInterface(
    fn=chat,
    title="tinychatbot-1b",
    description="tinychatbot running on llama.cpp + ZeroGPU",
)

demo.launch()