Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModel
+from transformers import pipeline
 import time
 from functools import wraps
 import sys
-import spaces  # Asegúrate de que este módulo esté disponible y correctamente instalado
+from multimolecule import RnaTokenizer, RnaFmModel  # Importar clases específicas de multimolecule
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -24,118 +24,86 @@ if device == "cpu":
     print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
 
 # Cargar el modelo y el tokenizador
-model_name = "yangheng/OmniGenome"
+model_name = "multimolecule/mrnafm"
 
 try:
     print("Cargando el tokenizador...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    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 = AutoModel.from_pretrained(model_name)
+    model = RnaFmModel.from_pretrained(model_name)
     model.to(device)
 except Exception as e:
     print(f"Error al cargar el modelo: {e}")
     sys.exit(1)
 
-@spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
 @medir_tiempo
-def predecir_estructura_rna(secuencias):
+def predecir_fill_mask(secuencias):
     """
-    Función que predice estructuras secundarias de ARN a partir de secuencias de ARN proporcionadas.
+    Función que realiza una predicción de Fill-Mask para las secuencias de ARN proporcionadas.
     """
     try:
         if not secuencias.strip():
-            return "Por favor, ingresa una o más secuencias de ARN válidas."
+            return "Por favor, ingresa una o más secuencias de ARN válidas con <mask> para predecir."
 
         # Separar las secuencias por líneas y eliminar espacios vacíos
         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:
-            # Validar la secuencia de ARN
-            if not all(residue in 'AUCG' for residue in seq):
-                resultados.append(f"Secuencia inválida: {seq}. Solo se permiten los nucleótidos A, U, C y G.")
+            # 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
 
-            # Tokenizar la secuencia
-            inputs = tokenizer(seq, return_tensors="pt")
-            input_ids = inputs["input_ids"].to(device)
-            attention_mask = inputs["attention_mask"].to(device)
-
-            # Aplicar el modelo para obtener los logits
-            with torch.no_grad():
-                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-
-            # Asumimos que el modelo devuelve logits para cada nucleótido que indican la estructura secundaria
-            # Debes ajustar esto según la arquitectura específica de OmniGenome
-
-            # Por ejemplo, supongamos que el modelo tiene una cabeza de clasificación con N etiquetas
-            # donde cada etiqueta representa una clase de estructura secundaria (e.g., Helix, Loop, etc.)
-
-            # Obtener las predicciones seleccionando la clase con el logit más alto
-            predictions = torch.argmax(outputs.logits, dim=-1).squeeze().tolist()
-
-            # Definir el mapeo de clases según la documentación del modelo OmniGenome
-            # Este mapeo debe ajustarse a las clases específicas que OmniGenome predice
-            # Por ejemplo:
-            structure_mapping = {
-                0: 'Helix',
-                1: 'Loop',
-                2: 'Bulge',
-                3: 'Internal Loop',
-                # Agrega más clases si es necesario
-            }
-
-            # Convertir las predicciones numéricas a etiquetas legibles
-            predicted_structures = [structure_mapping.get(pred, "Unknown") for pred in predictions]
-
-            # Emparejar cada nucleótido con su etiqueta de estructura predicha
-            nucleotide_to_structure = list(zip(list(seq), predicted_structures))
+            # Realizar la predicción de fill-mask
+            predictions = fill_mask(seq)
 
-            # Formatear el resultado para mostrarlo en la interfaz
-            secuencia_resultado = []
-            for i, (nucleotide, structure) in enumerate(nucleotide_to_structure):
-                secuencia_resultado.append(f"Posición {i+1} - {nucleotide}: {structure}")
+            # Formatear las predicciones
+            pred_str = ""
+            for pred in predictions:
+                pred_str += f"Predicción: {pred['sequence']}, Score: {pred['score']:.4f}\n"
 
-            # Unir las predicciones en un solo string
-            resultados.append("\n".join(secuencia_resultado))
+            resultados.append(f"Secuencia: {seq}\n{pred_str}")
 
-        # Unir los resultados de todas las secuencias separadas por dos saltos de línea
         return "\n\n".join(resultados)
 
     except Exception as e:
         print(f"Error durante la predicción: {e}")
-        return f"Error al predecir las estructuras de ARN: {e}"
+        return f"Error al realizar la predicción: {e}"
 
 # Definir la interfaz de Gradio
-titulo = "OmniGenome: Predicción de Estructuras Secundarias de ARN"
+titulo = "OmniGenome: Predicción de Fill-Mask para Secuencias de ARN"
 descripcion = (
-    "Ingresa una o más secuencias de ARN (una por línea) y obtén predicciones de estructuras secundarias para cada nucleótido."
-    " El modelo utilizado es OmniGenome, un modelo de fundamentos basado en transformadores para alineación secuencia-estructura en tareas genómicas."
+    "Ingresa una o más secuencias de ARN (una por línea) con un token <mask> donde deseas realizar la predicción."
+    " El modelo utilizado es mRNA-FM de MultiMolecule, un modelo pre-entrenado de lenguaje para secuencias de ARN."
 )
 
 iface = gr.Interface(
-    fn=predecir_estructura_rna,
+    fn=predecir_fill_mask,
     inputs=gr.Textbox(
         lines=10, 
-        placeholder="Escribe tus secuencias de ARN aquí, una por línea (solo A, U, C, G)...", 
-        label="Secuencias de ARN"
+        placeholder="Escribe tus secuencias de ARN aquí, una por línea, incluyendo <mask> donde desees predecir...", 
+        label="Secuencias de ARN con <mask>"
     ),
-    outputs=gr.Textbox(label="Predicciones de Estructuras Secundarias de ARN"),
+    outputs=gr.Textbox(label="Predicciones de Fill-Mask"),
     title=titulo,
     description=descripcion,
     examples=[
         [
-            "AUGGCUACUUUCG",
-            "GCGCGAUCGACGUAGCUAGC"
+            "AUGGCUACUUU<mask>G",
+            "GCGCGAU<mask>CGACGUAGCUAGC"
         ],
         [
-            "AUAUGCGGUAUCGUACGUA",
-            "GGAUACGUGAUCGUAGCAGU"
+            "AUAUGCGGUAUCGU<mask>GUA",
+            "GGAUACGUGAU<mask>GCUAGCAGU"
         ]
     ],
     cache_examples=False,