Update app.py

app.py

@@ -8,58 +8,50 @@ 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.")
-
-
-# Cargar el tokenizador
+# Configuración inicial
 tokenizer = AutoTokenizer.from_pretrained("Juliofc/chaterapia_model")
-
-# Añadir el token especial [PAD]
-#tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-
-# Cargar el modelo base y ajustar el tamaño de los embeddings de tokens
-model_base = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
+model_base = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it").to(device)
 model_base.resize_token_embeddings(len(tokenizer))
-
-# Cargar el modelo con el adaptador
 model_with_adapter = PeftModel.from_pretrained(model_base, "Juliofc/chaterapia_model").to(device)
 
-
-# Suponiendo que `tokenizer` y `model_with_adapter` ya están inicializados
-
-def chat_with_model(user_input, conversation_history=""):
-    # Asegúrate de que conversation_history tenga un valor inicial adecuado
-    if conversation_history is None:
-        conversation_history = ""
-
-    conversation_history += f"Usuario: {user_input}\n"
+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
+def generate_response(user_input, chat_history):
+    # Preparar el input agregando el historial de chat
+    mensajes = [{"content": user_input, "role": "user"}]
+    user_input  = tokenizer.apply_chat_template(mensajes, tokenize=False)
+    input_with_history = chat_history + user_input
     
-    input_tokens = tokenizer.encode(user_input, return_tensors='pt').to(device)
-    output_tokens = model_with_adapter.generate(input_tokens, max_new_tokens=50, pad_token_id=tokenizer.pad_token_id)
-    generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    input_tokens = tokenizer(input_with_history, return_tensors='pt', padding=True, truncation=True, max_length=1024).to(device)
     
-    conversation_history += f"Modelo: {generated_text}\n"
+    # 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)
     
-    return "", conversation_history
+    # Actualizar el historial del chat
+    new_chat_history = chat_history + "\n" + generated_text
+    return generated_text, new_chat_history
 
-# Define los componentes de la interfaz de Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("### Chat con IA")
-    input_text = gr.Textbox(label="Tu mensaje")
-    submit_button = gr.Button("Enviar")
-    output_text = gr.Textbox(label="Historial de la conversación", lines=10, interactive=False)
-    # Inicializa el estado conversation_history con una cadena vacía
-    conversation_history = gr.State(value="")
+# Gradio interface
+def chat_interface(user_input, chat_history=""):
+    response, new_chat_history = generate_response(user_input, chat_history)
+    return response, new_chat_history
 
-    submit_button.click(
-        fn=chat_with_model,
-        inputs=[input_text, conversation_history],
-        outputs=[input_text, conversation_history]
-    )
+iface = gr.Interface(fn=chat_interface,
+                     inputs=[gr.inputs.Textbox(lines=2, label="Tu Mensaje"), gr.inputs.Textbox(lines=10, label="Historial de Chat", default="")],
+                     outputs=[gr.outputs.Textbox(label="Respuesta del Modelo"), gr.outputs.Textbox(label="Nuevo Historial de Chat")],
+                     description="Chat con tu Modelo")
 
-# Asegúrate de lanzar usando un puerto disponible o soluciona el problema del puerto ocupado como se discutió anteriormente.
-demo.launch()
+iface.launch()