Update app.py

app.py

@@ -1,10 +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
-import spaces  # Asegúrate de que este módulo esté disponible y correctamente instalado
+import spaces  # Asegúrate de que este módulo esté correctamente instalado y disponible
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -23,73 +23,126 @@ 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 desde HuggingFace
+# Cargar el tokenizador y el modelo de generación desde HuggingFace
 model_name = "dmis-lab/selfbiorag_7b"
 
 try:
-    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)
-
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
+    print("Cargando el tokenizador y el modelo de generación desde HuggingFace...")
+    tokenizer_gen = AutoTokenizer.from_pretrained(model_name)
+    model_gen = AutoModelForCausalLM.from_pretrained(
+        model_name,
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     ).to(device)
+except ValueError as e:
+    print(f"Error al cargar el tokenizador de generación: {e}")
+    sys.exit(1)
 except Exception as e:
-    print(f"Error al cargar el modelo: {e}")
+    print(f"Error al cargar el modelo de generación: {e}")
     sys.exit(1)
 
+# Definir modelos de traducción
+# Diccionario de modelos de traducción según el idioma seleccionado
+translation_models = {
+    "Español": "Helsinki-NLP/opus-mt-en-es",
+    "Portugués": "Helsinki-NLP/opus-mt-en-pt",
+    "Francés": "Helsinki-NLP/opus-mt-en-fr",
+    "Alemán": "Helsinki-NLP/opus-mt-en-de",
+    # Añade más idiomas y sus respectivos modelos según sea necesario
+}
+
+# Cargar los tokenizadores y modelos de traducción
+tokenizer_trans = {}
+model_trans = {}
+
+for lang, model_name_trans in translation_models.items():
+    try:
+        print(f"Cargando el tokenizador y el modelo de traducción para {lang} desde HuggingFace...")
+        tokenizer = MarianTokenizer.from_pretrained(model_name_trans)
+        model = MarianMTModel.from_pretrained(model_name_trans).to(device)
+        tokenizer_trans[lang] = tokenizer
+        model_trans[lang] = model
+    except Exception as e:
+        print(f"Error al cargar el modelo de traducción para {lang}: {e}")
+        sys.exit(1)
+
 @spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
 @medir_tiempo
-def generar_respuesta(consulta):
+def generar_y_traducir_respuesta(consulta, idioma_destino):
     """
-    Función que genera una respuesta a partir de una consulta dada.
+    Función que genera una respuesta a partir de una consulta dada y la traduce al idioma seleccionado.
     """
     try:
         if not consulta.strip():
-            return "Por favor, ingresa una consulta válida."
+            return "Por favor, ingresa una consulta válida.", ""
         
         # Tokenizar la consulta
-        inputs = tokenizer.encode(consulta, return_tensors="pt").to(device)
+        inputs = tokenizer_gen.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       # Configura según la opción deseada
-            # "temperature": 0.0,    # Comenta o elimina si do_sample=False
-            # "top_p": 1.0            # Comenta o elimina si do_sample=False
+            "do_sample": False       # Generación determinista
+            # Puedes añadir otros parámetros como 'num_beams' si lo deseas
         }
         
         # Generar la respuesta
         with torch.no_grad():
-            outputs = model.generate(input_ids=inputs, **generation_kwargs)
+            outputs = model_gen.generate(input_ids=inputs, **generation_kwargs)
+        
+        # Decodificar la respuesta en inglés
+        respuesta_en = tokenizer_gen.decode(outputs[0], skip_special_tokens=True)
+        
+        # Traducir la respuesta al idioma seleccionado
+        if idioma_destino in translation_models:
+            tokenizer_tr = tokenizer_trans[idioma_destino]
+            model_tr = model_trans[idioma_destino]
+            
+            # Preparar la entrada para la traducción
+            traducir_inputs = tokenizer_tr.prepare_seq2seq_batch([respuesta_en], return_tensors="pt").to(device)
+            
+            # Realizar la traducción
+            with torch.no_grad():
+                traduccion_outputs = model_tr.generate(**traducir_inputs)
+            
+            # Decodificar la traducción
+            respuesta_traducida = tokenizer_tr.decode(traduccion_outputs[0], skip_special_tokens=True)
+        else:
+            respuesta_traducida = "Idioma de destino no soportado."
         
-        # Decodificar la respuesta
-        respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return respuesta
+        return respuesta_en, respuesta_traducida
     except Exception as e:
-        print(f"Error durante la generación de respuesta: {e}")
-        return f"Error al generar la respuesta: {e}"
+        print(f"Error durante la generación o traducción de la respuesta: {e}")
+        return f"Error al generar la respuesta: {e}", ""
 
 # Definir la interfaz de Gradio
-titulo = "Generador de Respuestas con SelfBioRAG 7B"
-descripcion = "Ingresa una consulta y el modelo generará una respuesta basada en el contenido proporcionado."
+titulo = "Generador y Traductor de Respuestas con SelfBioRAG 7B"
+descripcion = (
+    "Ingresa una consulta y el modelo generará una respuesta en inglés. "
+    "Luego, puedes seleccionar un idioma para traducir la respuesta generada."
+)
+
+idiomas_disponibles = list(translation_models.keys())
 
 iface = gr.Interface(
-    fn=generar_respuesta,
-    inputs=gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí..."),
-    outputs=gr.Textbox(),
+    fn=generar_y_traducir_respuesta,
+    inputs=[
+        gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí...", label="Consulta"),
+        gr.Dropdown(choices=idiomas_disponibles, value="Español", label="Idioma de Traducción")
+    ],
+    outputs=[
+        gr.Textbox(label="Respuesta en Inglés"),
+        gr.Textbox(label="Respuesta Traducida")
+    ],
     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."
+            "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.",
+            "Español"
         ],
         [
-            "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."
+            "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.",
+            "Portugués"
         ]
     ],
     cache_examples=False