import gradio as gr
from transformers import pipeline, login
import torch
from dotenv import load_dotenv
import os
import spaces  # Importar la librería de spaces

# Cargar variables de entorno desde el archivo .env
load_dotenv()

# Obtener el token de Hugging Face desde las variables de entorno
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")

# Verificar si el token fue cargado correctamente
if not huggingface_token:
    raise ValueError("Token de Hugging Face no encontrado en las variables de entorno")

# Iniciar sesión con tu token de Hugging Face
login(token=huggingface_token)  # Inicia sesión con el token

# Función para cargar el modelo en GPU (si está disponible) o en CPU si hay un error
def load_model():
    model_id = "facebook/MobileLLM-R1.5-950M"
    try:
        # Intentar cargar el modelo en GPU
        print("Intentando cargar el modelo en la GPU...")
        return pipeline(
            "text-generation",
            model=model_id,
            torch_dtype="auto",
            device_map="auto",  # Esto automáticamente usa la GPU si está disponible
        )
    except Exception as e:
        # Si hay un error (por ejemplo, no hay GPU disponible), usar CPU
        print("Error en GPU, cambiando a CPU. Error:", e)
        return pipeline(
            "text-generation",
            model=model_id,
            torch_dtype="auto",
            device=0,  # Usar CPU en caso de error
        )

# Cargar el modelo con el fallback a CPU
pipe = load_model()

# Decorador para garantizar que el código se ejecute en GPU si está disponible
@spaces.GPU
def chat_with_model(system_message, user_message):
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_message},
    ]
    
    # Generar la respuesta usando el modelo
    outputs = pipe(messages, max_new_tokens=8192)
    return outputs[0]["generated_text"][-1]

# Crear la interfaz de Gradio
gr.Interface(
    fn=chat_with_model,
    inputs=[
        gr.Textbox(label="System Message", placeholder="Escribe el mensaje del sistema..."),
        gr.Textbox(label="User Message", placeholder="Escribe tu mensaje aquí..."),
    ],
    outputs="text",
    live=True,
).launch()