from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import gradio as gr
import os
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

# Asegúrate de que tu token de Hugging Face está cargado como una variable de entorno
hf_token = os.environ.get("token")
if hf_token is not None:
    from huggingface_hub import HfFolder
    HfFolder.save_token(hf_token)

# Configuración inicial
tokenizer = AutoTokenizer.from_pretrained("somosnlp/chaterapia_model")
model_base = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it").to(device)
model_base.resize_token_embeddings(len(tokenizer))
model_with_adapter = PeftModel.from_pretrained(model_base, "somosnlp/chaterapia_model").to(device)

CHAT_TEMPLATE= """{% for message in messages %}
    {% if message['role'] == 'user' %}
        {{'<user> ' + message['content'].strip() + ' </user>' }}
    {% elif message['role'] == 'system' %}
        {{'<system>\\n' + message['content'].strip() + '\\n</system>\\n\\n' }}
    {% elif message['role'] == 'assistant' %}
        {{ message['content'].strip() + ' </assistant>' + eos_token }}
    {% elif message['role'] == 'input' %}
        {{'<input> ' + message['content'] + ' </input>' }}
    {% endif %}
{% endfor %}""" # Asegúrate de usar tu CHAT_TEMPLATE aquí
tokenizer.chat_template = CHAT_TEMPLATE

chat_history = []  # Historial de chat global
chatbot_text = []
def generate_response(user_input):
    global chat_history
    # Agregar input del usuario al historial
    chat_history.append({"content": user_input, "role": "user"})
    
    # Preparación del input para el modelo
    user_input = tokenizer.apply_chat_template(chat_history, tokenize=False)
    input_tokens = tokenizer(user_input, return_tensors='pt', padding=True, truncation=True, max_length=1024).to(device)
    
    # Generación de la respuesta del modelo
    output_tokens = model_with_adapter.generate(**input_tokens, max_length=1024, pad_token_id=tokenizer.eos_token_id, top_k=50, do_sample=True, top_p=0.95, temperature=0.7)
    generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
    
    # Extracción de la respuesta generada
    last_us = generated_text.rfind("</user>") + len("</user>")
    last_as = generated_text.rfind("</assistant>")
    generated_text = generated_text[last_us:last_as].strip()
    
    # Agregar la respuesta del bot al historial
    chat_history.append({"content": generated_text, "role": "assistant"})
    return generated_text

def respond(message):
    global chatbot_text
    if message:  # Verificar si el mensaje no está vacío
        bot_response = generate_response(message)
        chatbot_text.append((message, bot_response))
        return chatbot_text
    return [("", "")]


def clear_chat_and_history():
    global chat_history
    global chatbot_text
    chat_history.clear()# Vaciar el historial de chat
    chatbot_text.clear()
    return "", []

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    # Usar un Button regular en lugar de ClearButton para tener control sobre la función que se ejecuta
    with gr.Row():
        msg = gr.Textbox(label="Tu mensaje", placeholder="Escribe aquí...", lines=1)
        send_btn = gr.Button("Enviar")
    clear_btn = gr.Button("Limpiar Chat")
    
    # Acción al presionar el botón Enviar
    send_btn.click(fn=respond, inputs=msg, outputs=chatbot)
    clear_btn.click(fn=clear_chat_and_history, inputs=None, outputs=[msg, chatbot])

demo.launch()