Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import pipeline
 import time
 from functools import wraps
 import sys
-from multimolecule import RnaTokenizer, RnaFmModel  # Importar clases específicas de multimolecule
+import multimolecule  # Importar para registrar los modelos de multimolecule
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -19,26 +19,16 @@ def medir_tiempo(func):
     return wrapper
 
 # Verificar si CUDA está disponible para el modelo principal
-device = "cuda" if torch.cuda.is_available() else "cpu"
-if device == "cpu":
+device = 0 if torch.cuda.is_available() else -1
+if device == -1:
     print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
 
-# Cargar el modelo y el tokenizador
-model_name = "multimolecule/mrnafm"
-
+# Cargar el pipeline de Fill-Mask
 try:
-    print("Cargando el tokenizador...")
-    tokenizer = RnaTokenizer.from_pretrained(model_name)
-except ValueError as e:
-    print(f"Error al cargar el tokenizador: {e}")
-    sys.exit(1)
-
-try:
-    print("Cargando el modelo...")
-    model = RnaFmModel.from_pretrained(model_name)
-    model.to(device)
+    print("Cargando el pipeline de Fill-Mask...")
+    fill_mask = pipeline('fill-mask', model='multimolecule/mrnafm', device=device)
 except Exception as e:
-    print(f"Error al cargar el modelo: {e}")
+    print(f"Error al cargar el pipeline de Fill-Mask: {e}")
     sys.exit(1)
 
 @medir_tiempo
@@ -54,16 +44,13 @@ def predecir_fill_mask(secuencias):
         secuencias_lista = [seq.strip().upper() for seq in secuencias.strip().split('\n') if seq.strip()]
         resultados = []
 
-        # Crear el pipeline de fill-mask utilizando el tokenizador y modelo cargados
-        fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer, device=0 if device == "cuda" else -1)
-
         for seq in secuencias_lista:
             # Asegurarse de que la secuencia contenga al menos un <mask>
             if "<MASK>" not in seq and "<mask>" not in seq:
                 resultados.append(f"Secuencia sin token <mask>: {seq}. Agrega <mask> donde desees predecir.")
                 continue
 
-            # Realizar la predicción de fill-mask
+            # Realizar la predicción de Fill-Mask
             predictions = fill_mask(seq)
 
             # Formatear las predicciones