Update app.py

app.py

@@ -1,6 +1,11 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    MarianMTModel,
+    MarianTokenizer
+)
 import time
 from functools import wraps
 import sys
@@ -47,35 +52,21 @@ except Exception as 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 = {}
+# Definir el modelo de traducción al Español
+translation_model_name = "Helsinki-NLP/opus-mt-en-es"
 
-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, use_auth_token=hf_token)
-        model = MarianMTModel.from_pretrained(model_name_trans, use_auth_token=hf_token).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)
+try:
+    print(f"Cargando el tokenizador y el modelo de traducción para Español desde HuggingFace...")
+    tokenizer_tr_es = MarianTokenizer.from_pretrained(translation_model_name)
+    model_tr_es = MarianMTModel.from_pretrained(translation_model_name).to(device)
+except Exception as e:
+    print(f"Error al cargar el modelo de traducción para Español: {e}")
+    sys.exit(1)
 
 @medir_tiempo
 def generar_y_traducir_respuesta(consulta, idioma_destino):
     """
-    Función que genera una respuesta a partir de una consulta dada y la traduce al idioma seleccionado.
+    Función que genera una respuesta a partir de una consulta dada y la traduce al Español.
     """
     try:
         if not consulta.strip():
@@ -98,22 +89,15 @@ def generar_y_traducir_respuesta(consulta, idioma_destino):
         # Decodificar la respuesta en inglés con limpieza de espacios
         respuesta_en = tokenizer_gen.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=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.encode(respuesta_en, return_tensors="pt").to(device)
-            
-            # Realizar la traducción
-            with torch.no_grad():
-                traduccion_outputs = model_tr.generate(input_ids=traducir_inputs, max_length=512)
-            
-            # Decodificar la traducción con limpieza de espacios
-            respuesta_traducida = tokenizer_tr.decode(traduccion_outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        else:
-            respuesta_traducida = "Idioma de destino no soportado."
+        # Traducir la respuesta al Español
+        traducir_inputs = tokenizer_tr_es.encode(respuesta_en, return_tensors="pt").to(device)
+        
+        # Realizar la traducción
+        with torch.no_grad():
+            traduccion_outputs = model_tr_es.generate(input_ids=traducir_inputs, max_length=512)
+        
+        # Decodificar la traducción con limpieza de espacios
+        respuesta_traducida = tokenizer_tr_es.decode(traduccion_outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
         
         return respuesta_en, respuesta_traducida
     except Exception as e:
@@ -124,31 +108,26 @@ def generar_y_traducir_respuesta(consulta, idioma_destino):
 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."
+    "Luego, la respuesta se traducirá automáticamente al Español."
 )
 
-idiomas_disponibles = list(translation_models.keys())
-
 iface = gr.Interface(
     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")
+        gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí...", label="Consulta")
     ],
     outputs=[
         gr.Textbox(label="Respuesta en Inglés"),
-        gr.Textbox(label="Respuesta Traducida")
+        gr.Textbox(label="Respuesta Traducida al Español")
     ],
     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.",
-            "Español"
+            "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."
         ],
         [
-            "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"
+            "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