import gradio as gr
from huggingface_hub import InferenceClient

# Inicializar el cliente de inferencia
client = InferenceClient("Qwen/QwQ-32B-Preview")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # Construir el contexto de mensajes
    messages = [{"role": "system", "content": system_message}]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    
    response = ""
    
    # Realizar la llamada al modelo con streaming deshabilitado para capturar toda la respuesta
    full_response = client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=False,  # Deshabilitar streaming para capturar toda la respuesta de una vez
        temperature=temperature,
        top_p=top_p,
    )
    
    # Extraer el contenido de la respuesta
    if full_response and full_response.choices:
        response = full_response.choices[0].message['content']
    
    return response

# Definir los controles adicionales con mayor claridad y rangos adecuados
system_message_input = gr.Textbox(
    value="Responde siempre en español, y el código siempre completo",
    label="Mensaje del Sistema",
    lines=2,
    placeholder="Instrucciones para el modelo...",
)

max_tokens_slider = gr.Slider(
    minimum=512,
    maximum=8192,  # Aumentar el máximo si el modelo lo soporta
    step=1,
    value=2048,  # Aumentar el valor por defecto para respuestas más largas
    label="Máximo de Tokens Nuevos",
)

temperature_slider = gr.Slider(
    minimum=0.1,
    maximum=2.0,  # Ajustar el rango para mayor control
    step=0.1,
    value=0.7,
    label="Temperatura",
)

top_p_slider = gr.Slider(
    minimum=0.1,
    maximum=1.0,
    step=0.05,
    value=0.95,
    label="Top-p (muestreo por núcleo)",
)

# Configurar la interfaz de chat con una apariencia mejorada
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        system_message_input,
        max_tokens_slider,
        temperature_slider,
        top_p_slider,
    ],
    title="Interfaz de Chat con QwQ-32B-Preview",
    description="Interactúa con el modelo QwQ-32B-Preview de Hugging Face. Ajusta los parámetros para personalizar las respuestas.",
    theme="default",  # Puedes cambiar a otros temas como "huggingface", "dark", etc.
)

if __name__ == "__main__":
    demo.launch(share=True)  # `share=True` permite compartir la interfaz públicamente