app.py

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

# Modelo de lenguaje
model_id = "tiiuae/falcon-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Traductor inglés a español
translator = pipeline("translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es")

# Frases personalizadas (palabra clave : respuesta en español)
respuestas_personalizadas = {
    "últimamente me he sentido triste": "Lamento mucho que te sientas así. A veces hablar con alguien o hacer cosas que te gustan puede ayudar a sentirte un poco mejor. Estoy aquí para escucharte.",
    "siento que mis amigos me están excluyendo": "Es completamente válido sentirte triste cuando eso pasa. Las amistades pueden cambiar con el tiempo, pero eso no significa que no valgas mucho. Trata de hablar con ellos con calma o busca nuevos espacios donde te sientas valorado.",
    "a veces siento que no sirvo para nada": "Lamento que te sientas así. No eres un inútil, todos aprendemos a nuestro ritmo y eso está bien. Lo importante es que sigas esforzándote y que no te compares. Tú también puedes lograr cosas increíbles.",
    "le presté mi celular a mi amigo y lo rompió": "Entiendo que estés muy enojado, pero la violencia no es una buena opción. Intenta respirar profundo y hablar con él. Tal vez fue un accidente, y si no lo fue, siempre puedes pedir ayuda a un adulto para resolverlo.",
    "¿qué debo hacer?": "Lo siento, no puedo tomar decisiones por ti. Pero puedo ayudarte a reflexionar o darte información para que tú decidas lo mejor para ti.",
    "tomo esta medicina o no": "No puedo tomar decisiones médicas por ti. Siempre es mejor consultar a un médico o adulto responsable.",
    "mis papás quieren que me cambie de escuela, pero yo no quiero, ¿qué debería de hacer?": (
        "Entiendo que esta situación puede ser muy confusa y estresante. Cambiar de escuela es una decisión importante, "
        "y es normal que tengas muchas emociones al respecto.\n\n"
        "No puedo decirte qué decisión tomar, pero te animo a hablar con un adulto en quien confíes, como tus papás, un maestro o consejero escolar. "
        "Expresar lo que sientes puede ayudarte a que te escuchen y a tomar una decisión que se base en el diálogo y la comprensión. "
        "Estoy aquí para escucharte si quieres seguir hablando de lo que sientes."
    )
}

# Plantilla base
plantilla = """
You are a kind and empathetic assistant, specialized in supporting students with their mental health and emotional well-being. You always respond in Spanish, in a clear, natural, and compassionate tone. You validate the user's feelings, but you never encourage or justify violence. You also do not give direct advice on personal decisions, such as taking medicine or making life choices. If someone expresses a desire to hurt themselves or others, you help them calm down, reflect, and guide them toward healthy and safe ways to manage their emotions. Your responses should feel caring, non-judgmental, and supportive.

Example:
User: I’ve been feeling very alone lately.
Assistant: I’m really sorry to hear that. You're not alone, and sometimes sharing what you feel can help. Would you like to talk about it?

---
User: {entrada}
Assistant:
"""

# Función principal
def responder(entrada):
    entrada_lower = entrada.lower().strip()

    # Buscar coincidencia exacta con frases personalizadas
    for clave in respuestas_personalizadas:
        if clave in entrada_lower:
            return respuestas_personalizadas[clave]

    # Si no hay coincidencia, usar el modelo
    prompt = plantilla.format(entrada=entrada)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.85,
        top_p=0.92,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extraer la respuesta del asistente
    if "Assistant:" in response:
        respuesta_modelo = response.rsplit("Assistant:", 1)[-1].strip()
    else:
        respuesta_modelo = response.strip()

    # Cortar si incluye otra vez "User:", etc.
    for corte in ["User:", "Usuario:", "---"]:
        if corte in respuesta_modelo:
            respuesta_modelo = respuesta_modelo.split(corte)[0].strip()

    # Traducir al español
    traduccion = translator(respuesta_modelo, max_length=200)[0]['translation_text']
    return traduccion

# Interfaz Gradio
demo = gr.Interface(
    fn=responder,
    inputs=gr.Textbox(lines=5, label="¿Cómo te sientes o qué necesitas?"),
    outputs=gr.Textbox(label="Respuesta del Asistente (traducida al español)"),
    title="Asistente para Estudiantes con Falcon 7B + Traducción",
    description="Un asistente empático que responde en inglés y traduce al español. Apoya a estudiantes en temas emocionales sin tomar decisiones por ellos."
)

demo.launch()