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

REPO_ID = "unsloth/Kimi-K2.6-GGUF"
# Nota: Si el modelo está shardeado en varias partes (-00001-of-00014),
# llama.cpp requiere que descargues la primera parte y automáticamente detecta el resto,
# o debes usar un archivo único cuantizado menor (p. ej. Q2_K o un modelo de menor tamaño).
FILENAME = "UD-Q4_K_XL/Kimi-K2.6-UD-Q4_K_XL-00001-of-00014.gguf"

# Descargar archivo del Hub
model_path = hf_hub_download(
    repo_id=REPO_ID,
    filename=FILENAME
)

# Inicializar motor GGUF con soporte GPU
llm = Llama(
    model_path=model_path,
    n_gpu_layers=-1,       # Enviar capas a la GPU
    n_ctx=4096,
    verbose=False
)

@spaces.GPU(duration=120)
def generate_response(message, history, system_prompt=""):
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
            
    messages.append({"role": "user", "content": message})

    response_stream = llm.create_chat_completion(
        messages=messages,
        max_tokens=2048,
        temperature=0.6,
        top_p=0.9,
        stream=True
    )

    partial_text = ""
    for chunk in response_stream:
        delta = chunk["choices"][0]["delta"]
        if "content" in delta:
            partial_text += delta["content"]
            yield partial_text

demo = gr.ChatInterface(
    fn=generate_response,
    title="Kimi-K2.6 Inference",
    additional_inputs=[
        gr.Textbox("Eres un asistente experto en programación.", label="System Prompt")
    ]
)

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