Upload app.py

app.py

@@ -1,7 +1,7 @@
 """
-File: module_translation_MADLAD.py
+File: app.py
 
-Description: Module to translate between 400 languages.
+Description: Translate text...
 
 Author: Didier Guillevic
 Date: 2024-09-07
@@ -16,8 +16,8 @@ logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 import model_translation as translation
-from model_translation import tokenizer_multilingual
-from model_translation import model_multilingual
+from model_translation import tokenizer_multilingual, model_multilingual
+from model_translation import tokenizer_m2m100, model_m2m100
 
 from deep_translator import GoogleTranslator
 
@@ -116,7 +116,46 @@ def detect_language(text):
     lang = langdetect.detect(text)
     return lang
 
-@spaces.GPU
+def translate_with_bilingual_model(
+        text, src_lang, tgt_lang, sents_per_chunk
+    ):
+    """
+    Translate with Helsinki bilingual models
+    """
+    if src_lang not in translation.src_langs:
+        return (
+            f"ISSUE: currently no model for language '{src_lang}'. "
+             "If wrong language, please specify language."
+        )
+    logger.info(f"LANG: {src_lang}, TEXT: {text[:50]}...")
+    tokenizer, model = translation.get_tokenizer_model_for_src_lang(src_lang)
+    translated_text_bilingual_model = translate_with_model(
+        text, tokenizer, model, src_lang, sents_per_chunk)
+    return translated_text_bilingual_model
+
+
+#@spaces.GPU
+def translate_with_m2m100_model(
+        text: str,
+        src_lang: str,
+        tgt_lang: str,
+        sents_per_chunk: int=5):
+    """
+    Translate with the m2m100 model
+    """
+    tokenizer_m2m100.src_lang = src_lang
+    input_ids = tokenizer_m2m100(text, return_tensors="pt").input_ids.to(
+                model_m2m100.device)
+    outputs = model_m2m100.generate(
+        input_ids=input_ids,
+        forced_bos_token_id=tokenizer_m2m100.get_lang_id(tgt_lang)
+    )
+    translated_text = tokenizer_m2m100.batch_decode(
+        outputs[0], skip_special_tokens=True)
+    return translated_text
+
+
+#@spaces.GPU
 def translate_with_multilingual_model(
         text: str,
         tgt_lang: str,
@@ -124,7 +163,7 @@ def translate_with_multilingual_model(
         input_max_length: int=512,
         output_max_length: int=512):
     """
-    Translate the givent text into English (default "easy" language)
+    Translate the given text into English (default "easy" language)
     """
     chunks = build_text_chunks(text, None, sents_per_chunk)
     translated_chunks = []
@@ -139,7 +178,8 @@ def translate_with_multilingual_model(
                 model_multilingual.device)
         outputs = model_multilingual.generate(
             input_ids=input_ids, max_length=output_max_length)
-        translated_chunk = tokenizer_multilingual.decode(outputs[0], skip_special_tokens=True)
+        translated_chunk = tokenizer_multilingual.decode(
+            outputs[0], skip_special_tokens=True)
         translated_chunks.append(translated_chunk)
     
     return '\n'.join(translated_chunks)
@@ -153,25 +193,27 @@ def translate_text(
     """
     Translate the given text into English (default "easy" language)  
     """
+    src_lang = src_lang if (src_lang and src_lang != "auto") else detect_language(text)
+    tgt_lang = 'en' # Default "easy" language
+
     #
     # Bilingual (Helsinki model)
     #
-    src_lang = src_lang if (src_lang and src_lang != "auto") else detect_language(text)
-    if src_lang not in translation.src_langs:
-        return (
-            f"ISSUE: currently no model for language '{src_lang}'. "
-             "If wrong language, please specify language."
-        )
-    logger.info(f"LANG: {src_lang}, TEXT: {text[:50]}...")
-    tokenizer, model = translation.get_tokenizer_model_for_src_lang(src_lang)
+    translated_text_bilingual_model = translate_with_bilingual_model(
+        text, src_lang, tgt_lang, sents_per_chunk
+    )
+
+    #
+    # m2m100 model
+    #
+    translated_text_m2m100_model = translate_with_m2m100_model(
+        text, src_lang, tgt_lang, sents_per_chunk
+    )
 
-    translated_text_bilingual_model = translate_with_model(
-        text, tokenizer, model, src_lang, sents_per_chunk)
-    
     #
     # Multilingual model (Google MADLAD)
     #
-    tgt_lang = 'en' # Default "easy" language
+    
     translated_text_multilingual_model = translate_with_multilingual_model(
         text, tgt_lang, sents_per_chunk, input_max_length, output_max_length)
     
@@ -183,6 +225,7 @@ def translate_text(
 
     return (
         translated_text_bilingual_model,
+
         translated_text_multilingual_model,
         translated_text_google_translate
     )
@@ -207,6 +250,11 @@ with gr.Blocks() as demo:
         label="Bilingual translation model (Helsinki NLP)",
         render=False
     )
+    output_text_m2m100_model = gr.Textbox(
+        lines=6,
+        label="Facebook m2m100 translation model (**small**)",
+        render=False
+    )
     output_text_multilingual_model = gr.Textbox(
         lines=6,
         label="Multilingual translation model (**small** Google MADLAD)",
@@ -250,6 +298,7 @@ with gr.Blocks() as demo:
         outputs=[
             output_text_bilingual_model,
             output_text_multilingual_model,
+            output_text_m2m100_model,
             output_text_google_translate,
         ],
         additional_inputs=[sentences_per_chunk,],