chaterapia_demo / app2.py
Juliofc's picture
Rename app.py to app2.py
e69fc44 verified
raw
history blame
3.81 kB
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()