app2.py

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"
print(device)
# 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)
else:
    print("No se encontró el token de Hugging Face. Asegúrate de que la variable de entorno HF_TOKEN esté configurada.")

# Configuración inicial
tokenizer = AutoTokenizer.from_pretrained("Juliofc/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, "Juliofc/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

# Función para generar respuestas del modelo
import gradio as gr

# Asume que todas tus importaciones previas y configuraciones del modelo están aquí

# Aquí deberías tener definida la función `generate_response` tal como la compartiste
# Función para generar respuestas del modelo
def generate_response(user_input, chat_history):
    # Preparar el input agregando el historial de chat
    chat_history.append({"content": user_input, "role": "user"})
    
    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)
    
    # Generar la respuesta
    output_tokens = model_with_adapter.generate(**input_tokens, max_length=1024, pad_token_id=tokenizer.eos_token_id, top_k=50, top_p=0.95, temperature=0.7)
    generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)

    last_us = generated_text.rfind("</user>") + len("</user>")
    last_as = generated_text.rfind("</assistant>")
    generated_text = generated_text[last_us:last_as].strip() 
    chat_history.append({"content": generated_text, "role": "assistant"})
    return generated_text, chat_history
    
# Inicializa el historial de la conversación como una lista vacía
historial_de_chat = []

# Define la función que será llamada por la interfaz de Gradio
def chatbot_interface(pregunta):
    global historial_de_chat
    # Genera la respuesta utilizando tu modelo
    respuesta, historial_de_chat = generate_response(pregunta, historial_de_chat)
    # Prepara el historial de chat para mostrarlo en la interfaz
    historial_para_mostrar = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in historial_de_chat])
    return historial_para_mostrar

with gr.Blocks() as app:
    with gr.Row():
        gr.Markdown("### Chat de WhatsApp Simulado")
    with gr.Row():
        chatbox = gr.Textbox(label="Escribe tu pregunta aquí", placeholder="Hola, ¿cómo estás?")
    with gr.Row():
        boton_enviar = gr.Button("Enviar")
    historial = gr.Textbox(label="Historial de Chat", lines=20, interactive=False, placeholder="Aquí aparecerán tus mensajes y las respuestas.")

    boton_enviar.click(fn=chatbot_interface, inputs=chatbox, outputs=historial)

app.launch()