Create app.py

app.py

@@ -0,0 +1,67 @@
+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()