Update app.py

app.py

@@ -1,54 +1,44 @@
 import gradio as gr
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-import torch
+from transformers import MarianMTModel, MarianTokenizer
 
-# Dispositivo: GPU si está disponible, sino CPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Cargar modelo y tokenizer
-model_name = "facebook/mbart-large-50"
-tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
-model = MBartForConditionalGeneration.from_pretrained(model_name).to(device)
-
-# Lista de idiomas disponibles
-idiomas = {
-    "Inglés": "en_XX",
-    "Español": "es_XX",
-    "Francés": "fr_XX",
-    "Alemán": "de_XX",
-    "Ruso": "ru_RU",
-    "Chino": "zh_CN",
-    "Árabe": "ar_AR"
+# Diccionario de modelos multilenguaje OPUS-MT (origen -> destino)
+modelos = {
+    ("mul", "en"): "Helsinki-NLP/opus-mt-mul-en",
+    ("en", "es"): "Helsinki-NLP/opus-mt-en-es",
+    ("es", "en"): "Helsinki-NLP/opus-mt-es-en",
+    ("en", "fr"): "Helsinki-NLP/opus-mt-en-fr",
+    ("fr", "en"): "Helsinki-NLP/opus-mt-fr-en",
 }
 
+# Función para cargar modelo y tokenizer según idioma
+def cargar_modelo(origen, destino):
+    key = (origen, destino)
+    model_name = modelos.get(key, "Helsinki-NLP/opus-mt-mul-en")
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    model = MarianMTModel.from_pretrained(model_name)
+    return tokenizer, model
+
+# Función de traducción
 def traducir(texto, idioma_origen, idioma_destino):
-    """
-    Traduce texto de un idioma a otro usando mBART-50.
-    """
-    if not texto.strip():
-        return "Por favor ingresa algún texto."
-    
-    tokenizer.src_lang = idioma_origen
-    encoded = tokenizer(texto, return_tensors="pt").to(device)
-    generated_tokens = model.generate(
-        **encoded, 
-        forced_bos_token_id=tokenizer.lang_code_to_id[idioma_destino]
-    )
-    traduccion = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-    return traduccion
+    tokenizer, model = cargar_modelo(idioma_origen, idioma_destino)
+    inputs = tokenizer(texto, return_tensors="pt", padding=True)
+    translated = model.generate(**inputs)
+    return tokenizer.decode(translated[0], skip_special_tokens=True)
+
+# Lista de idiomas disponibles
+idiomas = ["mul", "en", "es", "fr"]
 
 # Crear interfaz Gradio
 iface = gr.Interface(
     fn=traducir,
     inputs=[
-        gr.Textbox(label="Texto a traducir", placeholder="Escribe aquí tu texto...", lines=4),
-        gr.Dropdown(label="Idioma de origen", choices=list(idiomas.values()), value="en_XX"),
-        gr.Dropdown(label="Idioma de destino", choices=list(idiomas.values()), value="es_XX")
+        gr.Textbox(label="Texto a traducir", lines=5, placeholder="Escribe aquí..."),
+        gr.Dropdown(idiomas, label="Idioma origen", value="mul"),
+        gr.Dropdown(idiomas, label="Idioma destino", value="en")
     ],
     outputs=gr.Textbox(label="Traducción"),
-    title="Traductor mBART-50",
-    description="Traduce texto entre múltiples idiomas usando mBART-50."
+    title="Traductor Multilenguaje OPUS-MT",
+    description="Traduce textos entre múltiples idiomas usando OPUS-MT en CPU."
 )
 
-# Ejecutar interfaz
 iface.launch()