app.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# --- Cargar el tokenizador y el modelo ---
# Se cargan una sola vez al inicio de la aplicación para mayor eficiencia.
# 'device_map="auto"' es crucial para que Hugging Face Spaces lo asigne correctamente
# a la GPU (si está disponible) o a la CPU.
# 'torch_dtype=torch.bfloat16' optimiza el uso de memoria en GPUs modernas.
tokenizer = AutoTokenizer.from_pretrained("tiiuae/Falcon-H1-1.5B-Deep-Instruct")
model = AutoModelForCausalLM.from_pretrained(
    "tiiuae/Falcon-H1-1.5B-Deep-Instruct",
    torch_dtype=torch.bfloat16, # Puedes probar con torch.float16 si da problemas
    device_map="auto"
)

# --- Definir la función de generación de texto ---
def generate_text(user_input, history):
    # Personalidad de Amside AI
    # Este es un "system prompt" que ayuda al modelo a mantener el personaje.
    # Se añade al inicio de cada conversación o antes de cada turno.
    # Ajusta esta parte para afinar la personalidad de Amside.
    system_message = (
        "Eres Amside AI, una inteligencia artificial creada por Hodelygil. "
        "Tu propósito principal es asistir en el estudio y el aprendizaje, "
        "proporcionando información y explicaciones detalladas. "
        "Sin embargo, también eres amigable y puedes mantener conversaciones informales y agradables. "
        "Responde de manera informativa y útil, pero con un tono conversacional."
    )

    # Formatear los mensajes para el modelo, incluyendo el historial de chat
    # El formato del "chat_template" es crucial para que el modelo entienda el diálogo.
    messages = [{"role": "system", "content": system_message}] # El system message va primero
    for human, bot in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": bot})
    messages.append({"role": "user", "content": user_input})

    # Codificar los mensajes en IDs de tokens
    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt")

    # Mover los IDs de tokens al mismo dispositivo donde se cargó el modelo (GPU/CPU)
    input_ids = input_ids.to(model.device)

    # Generar la respuesta del modelo
    outputs = model.generate(
        input_ids,
        max_new_tokens=500,  # Máximo de tokens para la respuesta del modelo
        do_sample=True,      # Activa el muestreo para respuestas más creativas
        temperature=0.7,     # Controla la aleatoriedad (menor = más predecible)
        top_p=0.9,           # Controla la diversidad de las palabras elegidas
        repetition_penalty=1.1 # Penaliza la repetición de frases
    )

    # Decodificar la respuesta y eliminar el prompt original para solo mostrar la respuesta de Amside
    # outputs[0] es la secuencia completa, input_ids.shape[1] es la longitud del prompt
    response = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)

    return response

# --- Configurar la interfaz de Gradio ---
# gr.ChatInterface es una forma sencilla de crear una interfaz de chatbot.
iface = gr.ChatInterface(
    fn=generate_text,
    chatbot=gr.Chatbot(height=400), # Altura de la ventana de chat
    textbox=gr.Textbox(placeholder="Escribe tu mensaje para Amside AI...", container=False, scale=7),
    title="Amside AI - Tu Compañero de Estudio y Conversación",
    description="Soy Amside AI, una inteligencia artificial creada por Hodelygil. Estoy aquí para ayudarte a estudiar y conversar.",
    examples=[
        "¿Qué es la fotosíntesis?",
        "Cuéntame un dato interesante.",
        "¿Cómo funciona la gravedad?",
        "Hola Amside, ¿cómo estás?",
        "Dame un resumen de la historia de España."
    ],
    theme="soft", # Un tema visual más suave para la interfaz
    retry_btn="Reintentar",
    undo_btn="Deshacer",
    clear_btn="Limpiar Historial"
)

# Lanzar la aplicación Gradio
iface.launch()