Update app.py

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
+from transformers import AutoTokenizer
 import time
 from functools import wraps
 import sys
 import spaces  # Asegúrate de que este módulo esté disponible y correctamente instalado
+from model import GenerRNA  # Importa tu modelo personalizado
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -18,125 +19,99 @@ def medir_tiempo(func):
         return resultado
     return wrapper
 
-# Verificar si CUDA está disponible para el modelo principal
+# Decorador para asignar GPU durante la ejecución de la función
+# Nota: Asegúrate de que el decorador @spaces.GPU esté disponible en tu entorno
+@spaces.GPU(duration=120)  # Asigna GPU durante 120 segundos
+@medir_tiempo
+def generar_rna_sequence(prompt, max_length=256):
+    """
+    Función que genera una secuencia de RNA a partir de una secuencia inicial dada.
+    """
+    try:
+        if not prompt.strip():
+            return "Por favor, ingresa una secuencia de inicio válida."
+
+        # Tokenizar la entrada
+        inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+        # Generar la secuencia
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_length=max_length,
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
+                temperature=0.7,
+                top_k=50,
+                top_p=0.95,
+                do_sample=True
+            )
+
+        # Decodificar la secuencia generada
+        generated_sequence = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return generated_sequence
+
+    except Exception as e:
+        print(f"Error durante la generación de secuencia: {e}")
+        return f"Error al generar la secuencia: {e}"
+
+# Configurar el dispositivo (GPU si 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 principal desde HuggingFace
-model_name = "dmis-lab/selfbiorag_7b"
-
+# Cargar el tokenizador
 try:
-    print("Cargando el tokenizador y el modelo desde HuggingFace...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    print("Cargando el tokenizador...")
+    tokenizer = AutoTokenizer.from_pretrained("tokenizer_bpe_1024")
 except ValueError as e:
     print(f"Error al cargar el tokenizador: {e}")
     sys.exit(1)
 
+# Cargar el modelo GenerRNA
 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)
-
-# 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
+    print("Cargando el modelo GenerRNA...")
+    model = GenerRNA()  # Instancia tu modelo personalizado
+    model.load_state_dict(torch.load("model.pt.recombined", map_location=device))
+    model.to(device)
+    model.eval()
+    print("Modelo GenerRNA cargado exitosamente.")
 except Exception as e:
-    print(f"Error al cargar el modelo de traducción: {e}")
+    print(f"Error al cargar el modelo GenerRNA: {e}")
     sys.exit(1)
 
-@spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
-@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)
-        
-        # 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}"
-
-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}"
-
-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
-
 # 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 = "GenerRNA - Generador de Secuencias de RNA"
+descripcion = (
+    "GenerRNA es un modelo generativo de RNA basado en una arquitectura Transformer. "
+    "Ingresa una secuencia inicial opcional y define la longitud máxima para generar nuevas secuencias de RNA."
+)
 
 iface = gr.Interface(
-    fn=procesar_consulta,
+    fn=generar_rna_sequence,
     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")
+        gr.Textbox(
+            lines=5,
+            placeholder="Ingresa una secuencia de RNA inicial (opcional)...",
+            label="Secuencia Inicial"
+        ),
+        gr.Slider(
+            minimum=50,
+            maximum=1000,
+            step=50,
+            value=256,
+            label="Longitud Máxima de la Secuencia"
+        )
     ],
+    outputs=gr.Textbox(label="Secuencia de RNA Generada"),
     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."
+            "AUGGCUACGUAUCGACGUA"
         ],
         [
-            "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."
+            "GCUAUGCUAGCUAGCUGAC"
         ]
     ],
     cache_examples=False,