Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
+from transformers import AutoTokenizer, EsmForTokenClassification
 import time
 from functools import wraps
 import sys
@@ -23,120 +23,103 @@ 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 principal desde HuggingFace
-model_name = "dmis-lab/selfbiorag_7b"
+# 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"
 
 try:
-    print("Cargando el tokenizador y el modelo desde HuggingFace...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    print("Cargando el tokenizador...")
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
 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)
+    print("Cargando el modelo de predicción...")
+    model = EsmForTokenClassification.from_pretrained(model_name)
+    model.to(device)
 except Exception as e:
     print(f"Error al cargar el modelo: {e}")
     sys.exit(1)
 
-# Cargar el modelo de traducción en CPU
-try:
-    print("Cargando el tokenizador y el modelo de traducción en CPU...")
-    translation_model_name = "Helsinki-NLP/opus-mt-en-es"
-    translator_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
-    translator_model = MarianMTModel.from_pretrained(translation_model_name).to("cpu")  # Forzar a CPU
-except Exception as e:
-    print(f"Error al cargar el modelo de traducción: {e}")
-    sys.exit(1)
-
 @spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
 @medir_tiempo
-def generar_respuesta(consulta):
+def predecir_sitios_arn(secuencias):
     """
-    Función que genera una respuesta a partir de una consulta dada.
+    Función que predice sitios de unión de ARN para las secuencias de proteínas proporcionadas.
     """
     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)
-        
-        # Configurar los parámetros de generación
-        generation_kwargs = {
-            "max_new_tokens": 100,  # Ajustado a 100
-            "do_sample": False       # No usar sampling
-            # "temperature": 0.6,    # Eliminado para evitar advertencias
-            # "top_p": 0.9            # Eliminado para evitar advertencias
-        }
-        
-        # Generar la respuesta
-        with torch.no_grad():
-            outputs = model.generate(input_ids=inputs, **generation_kwargs)
-        
-        # Decodificar la respuesta
-        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}"
+        if not secuencias.strip():
+            return "Por favor, ingresa una o más secuencias válidas."
 
-def traducir_texto(texto):
-    """
-    Función que traduce un texto de inglés a español.
-    """
-    try:
-        if not texto.strip():
-            return "No hay texto para traducir."
-        
-        # Tokenizar el texto a traducir
-        translated = translator_model.generate(**translator_tokenizer(texto, return_tensors="pt", padding=True))
-        
-        # Decodificar la traducción
-        traduccion = translator_tokenizer.decode(translated[0], skip_special_tokens=True)
-        return traduccion
-    except Exception as e:
-        print(f"Error durante la traducción: {e}")
-        return f"Error al traducir el texto: {e}"
+        # 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()]
+        resultados = []
 
-def procesar_consulta(consulta, idioma):
-    """
-    Función que procesa la consulta y devuelve la respuesta original y/o traducida según el idioma seleccionado.
-    """
-    respuesta_original = generar_respuesta(consulta)
-    
-    if idioma == "Español":
-        traduccion = traducir_texto(respuesta_original)
-    else:
-        traduccion = ""
-    
-    return respuesta_original, traduccion
+        for seq in secuencias_lista:
+            # Tokenizar la secuencia
+            inputs = tokenizer(seq, truncation=True, padding='max_length', max_length=1290, 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)
+
+            # 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]
+
+            # Emparejar cada residuo con su etiqueta predicha
+            residue_to_label = list(zip(list(seq), predicted_labels))
+
+            # 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}")
+            
+            resultados.append("\n".join(secuencia_resultado))
+
+        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}"
 
 # Definir la interfaz de Gradio
-titulo = "Generador de Respuestas con SelfBioRAG 7B"
-descripcion = "Ingresa una consulta y selecciona el idioma de salida. El modelo generará una respuesta basada en el contenido proporcionado."
+titulo = "ESM-2 para Predicción de Sitios de Unión 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."
+)
 
 iface = gr.Interface(
-    fn=procesar_consulta,
-    inputs=[
-        gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí...", label="Consulta"),
-        gr.Dropdown(choices=["Inglés", "Español"], value="Español", label="Idioma de Salida")
-    ],
-    outputs=[
-        gr.Textbox(label="Respuesta Original (Inglés)"),
-        gr.Textbox(label="Traducción al Español")
-    ],
+    fn=predecir_sitios_arn,
+    inputs=gr.Textbox(
+        lines=10, 
+        placeholder="Escribe tus secuencias de proteínas aquí, una por línea...", 
+        label="Secuencias de Proteínas"
+    ),
+    outputs=gr.Textbox(label="Predicciones de Sitios de Unión de ARN"),
     title=titulo,
     description=descripcion,
     examples=[
         [
-            "Clasifica el siguiente informe de radiología según la parte del cuerpo a la que se refiere (por ejemplo, pecho, abdomen, cerebro, etc.): Los discos intervertebrales en L4-L5 y L5-S1 muestran signos de degeneración con leve abultamiento que comprime la raíz nerviosa adyacente."
+            "VLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTK",
+            "SQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWF"
         ],
         [
-            "Resume los puntos clave sobre el papel de las mutaciones en los genes BRCA1 y BRCA2 en el aumento del riesgo de cáncer de mama."
+            "MKAILVVLLYTFATANADAVAHVAA",
+            "GATVQAAEEVTQGVVVVEEVAGGAA"
         ]
     ],
     cache_examples=False,