Update app.py

app.py

@@ -1,5 +1,9 @@
 import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer
+import torch
+
+# Limitar threads de CPU para optimizar
+torch.set_num_threads(4)
 
 # Diccionario de modelos multilenguaje OPUS-MT (origen -> destino)
 modelos = {
@@ -10,35 +14,59 @@ modelos = {
     ("fr", "en"): "Helsinki-NLP/opus-mt-fr-en",
 }
 
-# Función para cargar modelo y tokenizer según idioma
+# Cache de modelos cargados
+modelo_cache = {}
+
+# Función para cargar modelo y tokenizer según idioma (caché)
 def cargar_modelo(origen, destino):
     key = (origen, destino)
+    if key in modelo_cache:
+        return modelo_cache[key]
     model_name = modelos.get(key, "Helsinki-NLP/opus-mt-mul-en")
     tokenizer = MarianTokenizer.from_pretrained(model_name)
     model = MarianMTModel.from_pretrained(model_name)
+    modelo_cache[key] = (tokenizer, model)
     return tokenizer, model
 
-# Función de traducción
-def traducir(texto, idioma_origen, idioma_destino):
+# Dividir texto en fragmentos ≤512 tokens
+def dividir_texto(texto, max_len=512):
+    palabras = texto.split()
+    fragmentos = []
+    fragmento = []
+    for palabra in palabras:
+        fragmento.append(palabra)
+        if len(fragmento) >= max_len:
+            fragmentos.append(" ".join(fragmento))
+            fragmento = []
+    if fragmento:
+        fragmentos.append(" ".join(fragmento))
+    return fragmentos
+
+# Función de traducción optimizada para textos largos
+def traducir_fragmentos(texto, idioma_origen, idioma_destino):
     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)
+    fragmentos = dividir_texto(texto, max_len=512)
+    traduccion = []
+    for frag in fragmentos:
+        inputs = tokenizer(frag, return_tensors="pt", padding=True, truncation=True)
+        translated = model.generate(**inputs, max_length=512)
+        traduccion.append(tokenizer.decode(translated[0], skip_special_tokens=True))
+    return " ".join(traduccion)
 
 # Lista de idiomas disponibles
 idiomas = ["mul", "en", "es", "fr"]
 
 # Crear interfaz Gradio
 iface = gr.Interface(
-    fn=traducir,
+    fn=traducir_fragmentos,
     inputs=[
-        gr.Textbox(label="Texto a traducir", lines=5, placeholder="Escribe aquí..."),
+        gr.Textbox(label="Texto a traducir", lines=8, 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 Multilenguaje OPUS-MT",
-    description="Traduce textos entre múltiples idiomas usando OPUS-MT en CPU."
+    description="Traduce textos largos entre múltiples idiomas usando OPUS-MT en CPU. Optimizado para uso comercial."
 )
 
 iface.launch()