Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, EsmForTokenClassification
+from transformers import AutoTokenizer, AutoModel
 import time
 from functools import wraps
 import sys
@@ -23,25 +23,19 @@ 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.")
 
-# Definir el mapeo de clases
-class_mapping = {
-    0: 'Not Binding Site',
-    1: 'Binding Site',
-}
-
 # Cargar el modelo y el tokenizador
-model_name = "AmelieSchreiber/esm2_t6_8M_UR50D_rna_binding_site_predictor"
+model_name = "yangheng/OmniGenome"
 
 try:
     print("Cargando el tokenizador...")
-    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 except ValueError as e:
     print(f"Error al cargar el tokenizador: {e}")
     sys.exit(1)
 
 try:
-    print("Cargando el modelo de predicción...")
-    model = EsmForTokenClassification.from_pretrained(model_name)
+    print("Cargando el modelo...")
+    model = AutoModel.from_pretrained(model_name)
     model.to(device)
 except Exception as e:
     print(f"Error al cargar el modelo: {e}")
@@ -49,21 +43,26 @@ except Exception as e:
 
 @spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
 @medir_tiempo
-def predecir_sitios_arn(secuencias):
+def predecir_estructura_rna(secuencias):
     """
-    Función que predice sitios de unión de ARN para las secuencias de proteínas proporcionadas.
+    Función que predice estructuras secundarias de ARN a partir de secuencias de ARN proporcionadas.
     """
     try:
         if not secuencias.strip():
-            return "Por favor, ingresa una o más secuencias válidas."
+            return "Por favor, ingresa una o más secuencias de ARN válidas."
 
         # Separar las secuencias por líneas y eliminar espacios vacíos
-        secuencias_lista = [seq.strip() for seq in secuencias.strip().split('\n') if seq.strip()]
+        secuencias_lista = [seq.strip().upper() for seq in secuencias.strip().split('\n') if seq.strip()]
         resultados = []
 
         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.")
+                continue
+
             # Tokenizar la secuencia
-            inputs = tokenizer(seq, truncation=True, padding='max_length', max_length=1290, return_tensors="pt")
+            inputs = tokenizer(seq, return_tensors="pt")
             input_ids = inputs["input_ids"].to(device)
             attention_mask = inputs["attention_mask"].to(device)
 
@@ -71,55 +70,72 @@ def predecir_sitios_arn(secuencias):
             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()
 
-            # Convertir las predicciones a etiquetas
-            predicted_labels = [class_mapping.get(pred, "Unknown") for pred in predictions]
+            # 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 residuo con su etiqueta predicha
-            residue_to_label = list(zip(list(seq), predicted_labels))
+            # Emparejar cada nucleótido con su etiqueta de estructura predicha
+            nucleotide_to_structure = list(zip(list(seq), predicted_structures))
 
             # Formatear el resultado para mostrarlo en la interfaz
             secuencia_resultado = []
-            for i, (residue, label) in enumerate(residue_to_label):
-                # Omite los residuos 'PAD' que se agregan durante el padding
-                if residue != 'PAD':
-                    secuencia_resultado.append(f"Posición {i+1} - {residue}: {label}")
-            
+            for i, (nucleotide, structure) in enumerate(nucleotide_to_structure):
+                secuencia_resultado.append(f"Posición {i+1} - {nucleotide}: {structure}")
+
+            # Unir las predicciones en un solo string
             resultados.append("\n".join(secuencia_resultado))
 
+        # 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 los sitios de ARN: {e}"
+        return f"Error al predecir las estructuras de ARN: {e}"
 
 # Definir la interfaz de Gradio
-titulo = "ESM-2 para Predicción de Sitios de Unión de ARN"
+titulo = "OmniGenome: Predicción de Estructuras Secundarias de ARN"
 descripcion = (
-    "Ingresa una o más secuencias de proteínas (una por línea) y obtén predicciones de sitios de unión de ARN para cada residuo."
-    " El modelo utilizado es ESM-2, entrenado en el dataset 'S1' de sitios de unión proteína-ARN."
+    "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."
 )
 
 iface = gr.Interface(
-    fn=predecir_sitios_arn,
+    fn=predecir_estructura_rna,
     inputs=gr.Textbox(
         lines=10, 
-        placeholder="Escribe tus secuencias de proteínas aquí, una por línea...", 
-        label="Secuencias de Proteínas"
+        placeholder="Escribe tus secuencias de ARN aquí, una por línea (solo A, U, C, G)...", 
+        label="Secuencias de ARN"
     ),
-    outputs=gr.Textbox(label="Predicciones de Sitios de Unión de ARN"),
+    outputs=gr.Textbox(label="Predicciones de Estructuras Secundarias de ARN"),
     title=titulo,
     description=descripcion,
     examples=[
         [
-            "VLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTK",
-            "SQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWF"
+            "AUGGCUACUUUCG",
+            "GCGCGAUCGACGUAGCUAGC"
         ],
         [
-            "MKAILVVLLYTFATANADAVAHVAA",
-            "GATVQAAEEVTQGVVVVEEVAGGAA"
+            "AUAUGCGGUAUCGUACGUA",
+            "GGAUACGUGAUCGUAGCAGU"
         ]
     ],
     cache_examples=False,