Update app.py

app.py

@@ -1,21 +1,10 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
 import time
 from functools import wraps
 import sys
-
-# Intentar importar 'spaces' para usar el decorador GPU si está disponible
-try:
-    import spaces
-except ImportError:
-    # Si 'spaces' no está disponible, definir un decorador vacío
-    def GPU(duration):
-        def decorator(func):
-            return func
-        return decorator
-
-    spaces = type('spaces', (), {'GPU': GPU})
+import spaces  # Asegúrate de que este módulo esté disponible y correctamente instalado
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -29,110 +18,125 @@ def medir_tiempo(func):
         return resultado
     return wrapper
 
-# Configurar el dispositivo (GPU si está disponible)
+# Verificar si CUDA está disponible para el modelo principal
 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.")
 
-# Ruta local al tokenizador
-tokenizer_path = "tokenizer_bpe_1024"
+# Cargar el tokenizador y el modelo principal desde HuggingFace
+model_name = "dmis-lab/selfbiorag_7b"
 
-# Cargar el tokenizador desde el directorio local
 try:
-    print(f"Cargando el tokenizador desde el directorio local '{tokenizer_path}'...")
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    print("Cargando el tokenizador y el modelo desde HuggingFace...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 except ValueError as e:
     print(f"Error al cargar el tokenizador: {e}")
     sys.exit(1)
-except Exception as e:
-    print(f"Error inesperado al cargar el tokenizador: {e}")
-    sys.exit(1)
-
-# Ruta al modelo local
-model_path = "yangheng/OmniGenome"
 
-# Cargar el modelo desde el archivo local
 try:
-    print(f"Cargando el modelo GenerRNA desde '{model_path}'...")
     model = AutoModelForCausalLM.from_pretrained(
-        model_path,
+        model_name, 
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     ).to(device)
-    model.eval()
-    print("Modelo GenerRNA cargado exitosamente.")
-except FileNotFoundError:
-    print(f"Error: El archivo del modelo '{model_path}' no se encontró.")
+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 GenerRNA: {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_rna_sequence(prompt, max_length=256):
+def generar_respuesta(consulta):
     """
-    Función que genera una secuencia de RNA a partir de una secuencia inicial dada.
+    Función que genera una respuesta a partir de una consulta 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
+        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(
-                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
+            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 generación de secuencia: {e}")
-        return f"Error al generar la secuencia: {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 = "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."
-)
+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."
 
 iface = gr.Interface(
-    fn=generar_rna_sequence,
+    fn=procesar_consulta,
     inputs=[
-        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"
-        )
+        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")
     ],
-    outputs=gr.Textbox(label="Secuencia de RNA Generada"),
     title=titulo,
     description=descripcion,
     examples=[
         [
-            "AUGGCUACGUAUCGACGUA"
+            "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."
         ],
         [
-            "GCUAUGCUAGCUAGCUGAC"
+            "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."
         ]
     ],
     cache_examples=False,