Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import time
 from functools import wraps
 import sys
+import spaces  # Asegúrate de que este import es correcto según la documentación de Hugging Face
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -17,36 +18,17 @@ def medir_tiempo(func):
         return resultado
     return wrapper
 
-# Verificar si CUDA está disponible
-device = "cuda" if torch.cuda.is_available() else "cpu"
-if device == "cpu":
-    print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
-
-# Cargar el tokenizador y el modelo desde HuggingFace
-model_name = "dmis-lab/selfbiorag_7b"
-
-try:
-    print("Cargando el tokenizador y el modelo desde HuggingFace...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-except ValueError as e:
-    print(f"Error al cargar el tokenizador: {e}")
-    sys.exit(1)
-
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32
-    ).to(device)
-except Exception as e:
-    print(f"Error al cargar el modelo: {e}")
-    sys.exit(1)
-
+# Decorador para asignar GPU en Hugging Face Spaces
+@spaces.GPU(duration=120)
 @medir_tiempo
 def generar_respuesta(consulta):
     """
     Función que genera una respuesta a partir de una consulta dada.
     """
     try:
+        if not consulta.strip():
+            return "Por favor, ingresa una consulta válida."
+        
         # Tokenizar la consulta
         inputs = tokenizer.encode(consulta, return_tensors="pt").to(device)
         
@@ -67,8 +49,33 @@ def generar_respuesta(consulta):
         respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return respuesta
     except Exception as e:
+        print(f"Error durante la generación de respuesta: {e}")
         return f"Error al generar la respuesta: {e}"
 
+# Verificar si CUDA está disponible
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if device == "cpu":
+    print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
+
+# Cargar el tokenizador y el modelo desde HuggingFace
+model_name = "dmis-lab/selfbiorag_7b"
+
+try:
+    print("Cargando el tokenizador y el modelo desde HuggingFace...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+except ValueError as e:
+    print(f"Error al cargar el tokenizador: {e}")
+    sys.exit(1)
+
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    ).to(device)
+except Exception as e:
+    print(f"Error al cargar el modelo: {e}")
+    sys.exit(1)
+
 # Definir la interfaz de Gradio
 titulo = "Generador de Respuestas con SelfBioRAG 7B"
 descripcion = "Ingresa una consulta y el modelo generará una respuesta basada en el contenido proporcionado."