import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import spaces

# Cargar el modelo y el tokenizador
model_path = "Yhhxhfh/autotrain-2n9b3-mj6my"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",  # Uso de GPU si está disponible
    torch_dtype='auto'
).eval()

# Función para generar la respuesta en stream palabra por palabra
def chat_with_model_stream(user_input):
    # Convertir la conversación a tokens de entrada
    conversation = f"User: {user_input}\nAssistant:"
    input_ids = tokenizer(conversation, return_tensors='pt').input_ids
    
    # Generar el texto de forma incremental
    output_ids = model.generate(input_ids, max_length=10000, do_sample=True, pad_token_id=tokenizer.eos_token_id)
    
    # Decodificar palabra por palabra y devolver en stream
    response_text = ""
    for token_id in output_ids[0]:
        word = tokenizer.decode([token_id], skip_special_tokens=True)
        response_text += word + " "
        yield response_text.strip()

# Decorador para solicitar GPU en Hugging Face Spaces
@spaces.GPU(timeout=0)  # Solicitar GPU con tiempo de espera ilimitado
def create_interface():
    # Crear la interfaz personalizada con Blocks
    with gr.Blocks(css="""
/* Estilos CSS personalizados */
#chatbox {
    background-color: #f0f0f5;
    font-family: 'Arial', sans-serif;
    color: #333;
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0px 4px 12px rgba(0, 0, 0, 0.1);
    max-width: 700px;
    margin: auto;
}

#header {
    background-color: #4a90e2;
    color: white;
    padding: 20px;
    text-align: center;
    border-radius: 10px 10px 0 0;
}

#chat_area {
    min-height: 300px;
    margin: 10px 0;
    padding: 20px;
    background-color: #fff;
    border-radius: 8px;
    border: 1px solid #ddd;
}

#user_input {
    width: 100%;
    border-radius: 8px;
    border: 1px solid #ddd;
    padding: 10px;
    font-size: 16px;
}

#submit_button {
    background-color: #4a90e2;
    color: white;
    border-radius: 8px;
    padding: 10px 20px;
    cursor: pointer;
    border: none;
    margin-top: 10px;
}

""") as demo:

        with gr.Column(elem_id="chatbox"):  # Cambié gr.Box a gr.Column
            # Encabezado
            gr.Markdown(
                "<div id='header'><h1>Professional AI Chatbot</h1><p>Chat with our AI assistant</p></div>"
            )

            # Área de chat
            chat_area = gr.Textbox(label="Conversation", interactive=False, elem_id="chat_area")

            # Entrada de usuario
            user_input = gr.Textbox(placeholder="Enter your message here...", label="Your message", elem_id="user_input")

            # Botón para enviar el mensaje
            submit_button = gr.Button("Send", elem_id="submit_button")

            # Evento de envío (streaming de respuesta)
            submit_button.click(fn=chat_with_model_stream, inputs=user_input, outputs=chat_area)

    # Ejecutar la interfaz
    demo.launch(enable_queue=False)

# Llamar la función para crear la interfaz
create_interface()