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

# CONFIGURACIÓN OPTIMIZADA
REPO_ID = "Qwen/Qwen2.5-3B-Instruct-GGUF"
FILENAME = "qwen2.5-3b-instruct-q4_k_m.gguf" # Sugerencia: Q4_K_M es un 20% más rápido que Q5

print(f"Descargando {FILENAME}...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

print("Cargando modelo con optimizaciones de CPU...")
llm = Llama(
    model_path=model_path,
    n_ctx=1024,        # Reducido de 4096 a 1024. Esto acelera drásticamente el inicio.
    n_threads=4,       # Mantenemos 4 para aprovechar el Space
    n_batch=128,       # Reducido de 512 a 128 para procesar bloques más ligeros
    use_mlock=True,    # Intenta mantener el modelo en RAM física
    low_vram=True      # Optimiza el uso de memoria
)

def predict(message, system_prompt="Responde en español de forma breve."):
    # ChatML format
    prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    # max_tokens=512 es suficiente para la mayoría de respuestas y termina antes
    output = llm(prompt, max_tokens=512, stop=["<|im_end|>", "<|im_start|>"], echo=False)
    return output['choices'][0]['text']

with gr.Blocks() as demo:
    gr.Markdown(f"# Nodo API Optimizado: {REPO_ID}")
    with gr.Row():
        msg = gr.Textbox(label="Input")
        sys = gr.Textbox(label="System Prompt", value="Responde de forma concisa.")
    out = gr.Textbox(label="Output")
    btn = gr.Button("Generar")
    
    btn.click(predict, [msg, sys], out, api_name="query")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)