app.py

import torch
from transformers import pipeline
import gradio as gr
import time

# Cargar el modelo al iniciar
def load_model():
    print("🚀 Cargando Mistral-7B-Grok...")
    model = pipeline(
        "text-generation",
        model="HuggingFaceH4/mistral-7b-grok",
        torch_dtype=torch.float16,
        device_map="auto",
        max_new_tokens=256,
        temperature=0.7,
        do_sample=True
    )
    print("✅ Modelo cargado exitosamente!")
    return model

# Cargar modelo
model = load_model()

# Función para generar respuestas educativas
def cosmobot_response(message, history):
    """Genera respuesta usando Mistral-7B con contexto educativo"""
    
    # Crear prompt educativo
    system_prompt = """Eres CosmoBot, un asistente educativo especializado en exploración espacial, astronomía y cosmología. 

Tu misión es proporcionar explicaciones claras, precisas y accesibles para estudiantes de todos los niveles. Mantén un tono amigable y entusiasta sobre el cosmos.

Áreas de expertise:
- Planetas y sistema solar
- Estrellas y astrofísica  
- Galaxias y cosmología
- Exploración espacial y misiones
- Agujeros negros y fenómenos cósmicos
- Historia de la astronomía

Responde de manera educativa e inspiradora:"""

    # Construir conversación
    conversation = system_prompt + "\n\n"
    
    # Agregar historial de conversación
    for user_msg, bot_msg in history:
        conversation += f"Usuario: {user_msg}\nCosmoBot: {bot_msg}\n"
    
    # Agregar mensaje actual
    conversation += f"Usuario: {message}\nCosmoBot:"
    
    try:
        # Generar respuesta
        start_time = time.time()
        response = model(
            conversation,
            max_new_tokens=300,
            temperature=0.7,
            do_sample=True,
            top_p=0.9
        )[0]['generated_text']
        
        # Extraer solo la última respuesta
        bot_response = response.split("CosmoBot:")[-1].strip()
        
        # Limpiar respuesta
        if "Usuario:" in bot_response:
            bot_response = bot_response.split("Usuario:")[0].strip()
        
        response_time = time.time() - start_time
        print(f"⏱️ Tiempo de respuesta: {response_time:.2f}s")
        
        return bot_response
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return "🔧 Error temporal. Por favor, intenta nuevamente."

# Crear interfaz Gradio
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="🚀 CosmoBot - Asistente Espacial con Mistral-7B"
) as demo:
    
    gr.Markdown(
        """
        # 🚀 CosmoBot 
        ### Asistente Educativo Espacial con Mistral-7B-Grok
        
        ¡Hola! Soy **CosmoBot**, tu guía personal para explorar el universo. 
        Puedo explicarte sobre planetas, estrellas, galaxias, exploración espacial y más.
        
        **Ejemplos de preguntas:**
        - "Explícame qué son los agujeros negros"
        - "Háblame sobre la exploración de Marte" 
        - "Qué es la materia oscura?"
        - "Cuéntame sobre las misiones de la NASA"
        """
    )
    
    # Chatbot interactivo
    chatbot = gr.Chatbot(
        label="Conversación con CosmoBot",
        height=500,
        show_copy_button=True
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Escribe tu pregunta sobre el espacio",
            placeholder="Pregunta sobre astronomía, planetas, estrellas...",
            scale=4
        )
        submit = gr.Button("🚀 Enviar", scale=1)
    
    # Ejemplos rápidos
    gr.Examples(
        examples=[
            "Explícame qué son los agujeros negros",
            "Háblame sobre la misión Artemis de la NASA", 
            "Cómo se forman las estrellas?",
            "Qué sabemos sobre el planeta Marte?",
            "Explica la teoría del Big Bang"
        ],
        inputs=msg
    )
    
    # Funcionalidad de limpiar
    clear = gr.Button("🧹 Limpiar conversación")
    
    # Manejar envío de mensajes
    def respond(message, chat_history):
        if not message.strip():
            return "", chat_history
        
        bot_message = cosmobot_response(message, chat_history)
        chat_history.append((message, bot_message))
        return "", chat_history
    
    # Conectar eventos
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    submit.click(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

# Lanzar aplicación
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        share=False
    )