Initial commit: bilingual models, multilingual mode, Google Translate

app.py

@@ -8,58 +8,185 @@ Date: 2024-09-07
 """
 import spaces
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from transformers import BitsAndBytesConfig
 import gradio as gr
+import langdetect
 
-#
-# Load the "small" MADLAD400 model
-#
-model_name = "google/madlad400-10b-mt"
-
-#quantization_config = BitsAndBytesConfig(
-#    load_in_4bit=True,
-#    bnb_4bit_compute_dtype=torch.float16,
-#    bnb_4bit_use_double_quant=True,
-#    bnb_4bit_quant_type="nf4"
-#)
-quantization_config = BitsAndBytesConfig(
-    load_in_8bit=True,
-    llm_int8_threshold=200.0 # https://discuss.huggingface.co/t/correct-usage-of-bitsandbytesconfig/33809/5
-)
-
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    torch_dtype=torch.float16,
-    quantization_config=quantization_config)
-model = torch.compile(model)
+import logging
+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 deep_translator import GoogleTranslator
+
+from model_spacy import nlp_xx
 
 #
-# Translate given input text in given target language
+# Translate given input text
 #
+def build_text_chunks(text, src_lang, sents_per_chunk):
+    """
+    Given a text:
+    - Split the text into sentences.
+    - Build text chunks:
+        - Consider up to sents_per_chunk
+        - Ensure that we do not exceed translation.max_words_per_chunk
+    """
+    # Split text into sentences...
+    sentences = [
+        sent.text.strip() for sent in nlp_xx(text).sents if sent.text.strip()]
+    logger.info(f"LANG: {src_lang}, TEXT: {text[:20]}, NB_SENTS: {len(sentences)}")
+    
+    # Create text chunks of N sentences
+    chunks = []
+    chunk = ''
+    chunk_nb_sentences = 0
+    chunk_nb_words = 0
+
+    for i in range(0, len(sentences)):
+        # Get sentence
+        sent = sentences[i]
+        sent_nb_words = len(sent.split())
+
+        # If chunk already 'full', save chunk, start new chunk
+        if (
+                (chunk_nb_words + sent_nb_words > translation.max_words_per_chunk) or
+                (chunk_nb_sentences + 1 > sents_per_chunk)
+           ):
+            chunks.append(chunk)
+            chunk = ''
+            chunk_nb_sentences = 0
+            chunk_nb_words = 0
+        
+        # Append sentence to current chunk. One sentence per line.
+        chunk = (chunk + '\n' + sent) if chunk else sent
+        chunk_nb_sentences += 1
+        chunk_nb_words += sent_nb_words
+
+    # Append last chunk
+    if chunk:
+        chunks.append(chunk)
+    
+    # !!! SKIP splitting of text into chunks for now !!!
+    # Might not be reliable for non-European languages.
+    #chunks = [text, ]
+
+    # NOTE: The 'fa' (Persian) model has multiple target languages to choose from.
+    # We need to specifiy the desired languages among: fra ita por ron spa
+    #   https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-fa-itc
+    # Prepend text with >>fra<< in order to translate in French.
+    if src_lang == 'fa':
+        chunks = [">>fra<< " + chunk for chunk in chunks]
+
+    return chunks
+
+def translate_with_model(
+        text, tokenizer, model, src_lang, sents_per_chunk,
+        input_max_length=512, output_max_length=512):
+    
+    # Build text chunks (using sents_per_chunk and translation.max_words_per_chunk)
+    chunks = build_text_chunks(text, src_lang, sents_per_chunk)
+    logger.info(f"LANG: {src_lang}, TEXT: {text[:20]}, NB_CHUNKS: {len(chunks)}")
+
+    # Translate chunks
+    translated_chunks = []
+    for chunk in chunks:
+        inputs = tokenizer(
+            chunk, return_tensors="pt",
+            max_length=input_max_length,
+            truncation=True, padding="longest").to(model.device)
+        
+        outputs = model.generate(
+            **inputs,
+            max_length=output_max_length)
+        
+        translated_chunk = tokenizer.batch_decode(
+            outputs, skip_special_tokens=True)[0]
+        
+        #logger.info(f"Text: {chunk}")
+        #logger.info(f"Translation: {translated_chunk}")
+        
+        translated_chunks.append(translated_chunk)
+
+    return '\n'.join(translated_chunks)
+
+def detect_language(text):
+    lang = langdetect.detect(text)
+    return lang
+
 @spaces.GPU
-def translate_text(
+def translate_with_multilingual_model(
         text: str,
         tgt_lang: str,
-        input_max_length=512,
-        output_max_length=512) -> str:
+        sents_per_chunk: int=5,
+        input_max_length: int=512,
+        output_max_length: int=512):
     """
-    Translate text in given target language.
-    Input text will be split into chunk that will be translated sequentially.
-    We will have up to sents_per_chunk sentences in a given chunk.
+    Translate the givent text into English (default "easy" language)
     """
-    if not tgt_lang:
-        tgt_lang = "en"
-    input_text = f"<2{tgt_lang}> {text}"
-    input_ids = tokenizer(
-        input_text, return_tensors="pt",
-        max_length=input_max_length,
-        truncation=True, padding="longest").input_ids.to(model.device)
-    outputs = model.generate(input_ids=input_ids, max_length=output_max_length)
-    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return translated_text
+    chunks = build_text_chunks(text, None, sents_per_chunk)
+    translated_chunks = []
+
+    for chunk in chunks:
+        input_text = f"<2{tgt_lang}> {text}"
+        logger.info(f" Translating: {input_text[:30]}")
+        input_ids = tokenizer_multilingual(
+            input_text, return_tensors="pt",
+            max_length=input_max_length,
+            truncation=True, padding="longest").input_ids.to(
+                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_chunks.append(translated_chunk)
+    
+    return '\n'.join(translated_chunks)
+
+def translate_text(
+        text: str,
+        src_lang: str=None,
+        sents_per_chunk: int=5,
+        input_max_length: int=512,
+        output_max_length: int=512):
+    """
+    Translate the given text into English (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_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)
+    
+    #
+    # Google Translate
+    #
+    translated_text_google_translate = GoogleTranslator(
+        source='auto', target='en').translate(text=text)
+
+    return (
+        translated_text_bilingual_model,
+        translated_text_multilingual_model,
+        translated_text_google_translate
+    )
+
 
 #
 # User interface
@@ -67,32 +194,66 @@ def translate_text(
 with gr.Blocks() as demo:
 
     gr.Markdown("""
-        ## Text translation (basic, small paragraph, multilingual)
+        ## Text translation v0.0.2 (basic, small paragraph, multilingual)
     """)
     input_text = gr.Textbox(
-        lines=5,
+        lines=15,
         placeholder="Enter text to translate",
         label="Text to translate",
         render=False
     )
-    output_text = gr.Textbox(lines=10, render=False)
-    tgt_lang = gr.Radio(
-        choices=["en", "fr"], value="en", label="Target language", render=False)
+    output_text_bilingual_model = gr.Textbox(
+        lines=6,
+        label="Bilingual translation model (Helsinki NLP)",
+        render=False
+    )
+    output_text_multilingual_model = gr.Textbox(
+        lines=6,
+        label="Multilingual translation model (**small** Google MADLAD)",
+        render=False
+    )
+    output_text_google_translate = gr.Textbox(
+        lines=6,
+        label="Google Translate",
+        render=False
+    )
+
+    # Extra (additional) input parameters
+    sentences_per_chunk = gr.Slider(
+        minimum=1, maximum=10, value=5, step=1,
+        label="nb sentences per context",
+        render=False
+    )
+    src_lang = gr.Radio(
+        choices=["auto", "ar", "en", "fa", "fr", "he", "ja", "zh"], value="auto",
+        label="Source language",
+        render=False
+    )
 
     # Examples
     examples = [
-        ["ریچارد مور، رئیس سازمان مخفی اطلاعاتی بریتانیا (ام‌آی‌۶) در دیدار ویلیام برنز، رئیس سازمان اطلاعات مرکزی آمریکا (سیا) گفت همچنان احتمال اقدام ایران علیه اسرائیل در واکنش به ترور اسماعیل هنیه، رهبر حماس وجود دارد. آقای برنز نیز در این دیدار فاش کرد که در سال اول جنگ اوکراین، «خطر واقعی» وجود داشت که روسیه به استفاده از «تسلیحات هسته‌ای تاکتیکی» متوسل شود. این دو مقام امنیتی هشدار دادند که «نظم جهانی» از زمان جنگ سرد تا کنون تا این حد «در معرض تهدید» نبوده است.", "en"],
-        ["Clément Delangue est, avec Julien Chaumond et Thomas Wolf, l’un des trois Français cofondateurs de Hugging Face, une start-up d’intelligence artificielle (IA) de premier plan. Valorisée à 4,2 milliards d’euros après avoir levé près de 450 millions d’euros depuis sa création en 2016, cette société de droit américain est connue comme la plate-forme de référence où développeurs et entreprises publient des outils et des modèles pour faire de l’IA en open source, c’est-à-dire accessible gratuitement et modifiable.", "en"],
-        ["يُعد تفشي مرض جدري القردة قضية صحية عالمية خطيرة، ومن المهم محاولة منع انتشاره للحفاظ على سلامة الناس وتجنب العدوى. د. صموئيل بولاند، مدير الحوادث الخاصة بمرض الجدري في المكتب الإقليمي لمنظمة الصحة العالمية في أفريقيا، يتحدث من كينشاسا في جمهورية الكونغو الديمقراطية، ولديه بعض النصائح البسيطة التي يمكن للناس اتباعها لتقليل خطر انتشار المرض.", "en"],
-        ["【ワシントン＝冨山優介】米ボーイングの新型宇宙船「スターライナー」は７日午前０時（日本時間７日午後１時）過ぎ、米ニューメキシコ州のホワイトサンズ宇宙港に着地し、地球に帰還した。スターライナーは米宇宙飛行士２人を乗せて６月に打ち上げられ、国際宇宙ステーション（ＩＳＳ）に接続したが、機体のトラブルが解決できず、無人でＩＳＳから離脱した。", "en"],
-        ["張先生稱，奇瑞已經凖備在西班牙生產汽車，並決心採取「本地化」的方式進入歐洲市場。此外，他也否認該公司的出口受益於不公平補貼。奇瑞成立於1997年，是中國最大的汽車公司之一。它已經是中國最大的汽車出口商，並且制定了進一步擴張的野心勃勃的計劃。", "en"],
-        ["ברוכה הבאה, קיטי: בית הקפה החדש בלוס אנג'לס החתולה האהובה והחברים שלה מקבלים בית קפה משלהם בשדרות יוניברסל סיטי, שם תוכלו למצוא מגוון של פינוקים מתוקים – החל ממשקאות ועד עוגות", "en"],
+        ["ریچارد مور، رئیس سازمان مخفی اطلاعاتی بریتانیا (ام‌آی‌۶) در دیدار ویلیام برنز، رئیس سازمان اطلاعات مرکزی آمریکا (سیا) گفت همچنان احتمال اقدام ایران علیه اسرائیل در واکنش به ترور اسماعیل هنیه، رهبر حماس وجود دارد. آقای برنز نیز در این دیدار فاش کرد که در سال اول جنگ اوکراین، «خطر واقعی» وجود داشت که روسیه به استفاده از «تسلیحات هسته‌ای تاکتیکی» متوسل شود. این دو مقام امنیتی هشدار دادند که «نظم جهانی» از زمان جنگ سرد تا کنون تا این حد «در معرض تهدید» نبوده است.", "fa"],
+        ["Clément Delangue est, avec Julien Chaumond et Thomas Wolf, l’un des trois Français cofondateurs de Hugging Face, une start-up d’intelligence artificielle (IA) de premier plan. Valorisée à 4,2 milliards d’euros après avoir levé près de 450 millions d’euros depuis sa création en 2016, cette société de droit américain est connue comme la plate-forme de référence où développeurs et entreprises publient des outils et des modèles pour faire de l’IA en open source, c’est-à-dire accessible gratuitement et modifiable.", "fr"],
+        ["يُعد تفشي مرض جدري القردة قضية صحية عالمية خطيرة، ومن المهم محاولة منع انتشاره للحفاظ على سلامة الناس وتجنب العدوى. د. صموئيل بولاند، مدير الحوادث الخاصة بمرض الجدري في المكتب الإقليمي لمنظمة الصحة العالمية في أفريقيا، يتحدث من كينشاسا في جمهورية الكونغو الديمقراطية، ولديه بعض النصائح البسيطة التي يمكن للناس اتباعها لتقليل خطر انتشار المرض.", "ar"],
+        ["【ワシントン＝冨山優介】米ボーイングの新型宇宙船「スターライナー」は７日午前０時（日本時間７日午後１時）過ぎ、米ニューメキシコ州のホワイトサンズ宇宙港に着地し、地球に帰還した。スターライナーは米宇宙飛行士２人を乗せて６月に打ち上げられ、国際宇宙ステーション（ＩＳＳ）に接続したが、機体のトラブルが解決できず、無人でＩＳＳから離脱した。", "ja"],
+        ["張先生稱，奇瑞已經凖備在西班牙生產汽車，並決心採取「本地化」的方式進入歐洲市場。此外，他也否認該公司的出口受益於不公平補貼。奇瑞成立於1997年，是中國最大的汽車公司之一。它已經是中國最大的汽車出口商，並且制定了進一步擴張的野心勃勃的計劃。", "zh"],
+        ["ברוכה הבאה, קיטי: בית הקפה החדש בלוס אנג'לס החתולה האהובה והחברים שלה מקבלים בית קפה משלהם בשדרות יוניברסל סיטי, שם תוכלו למצוא מגוון של פינוקים מתוקים – החל ממשקאות ועד עוגות", "he"],
     ]
+
+    outputs = gr.Row(
+        
+    )
     
     gr.Interface(
         fn=translate_text,
-        inputs=[input_text, tgt_lang,],
-        outputs=[output_text,],
+        inputs=[input_text,],
+        outputs=[
+            output_text_bilingual_model,
+            output_text_multilingual_model,
+            output_text_google_translate,
+        ],
+        additional_inputs=[src_lang, sentences_per_chunk],
+        clear_btn=None,  # Unfortunately, clear_btn also reset the additional inputs. Hence disabling for now.
         allow_flagging="never",
         examples=examples,
         cache_examples=True
@@ -100,9 +261,10 @@ with gr.Blocks() as demo:
 
     with gr.Accordion("Documentation", open=False):
         gr.Markdown("""
-            - Model: Serving the smallest (3b) [MADLAD-400 model](https://arxiv.org/abs/2309.04662)
-            - Basic: does not process long paragraph / document
+            - Models: serving bilingual models from Helsinki NLP and multilingual model from Google MADLAD.
+            - Basic: processing of long paragraph / document to be enhanced.
             - Most examples are copy/pasted from BBC news international web sites.
         """)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

model_spacy.py

@@ -0,0 +1,24 @@
+"""
+File: model_spacy.py
+
+Description:
+    Load a spaCy model (will be used to split a text into sentences)
+
+Author: Didier Guillevic
+Date: 2024-03-30
+"""
+
+import spacy
+
+model_xx_name = 'xx_sent_ud_sm'
+
+nlp_xx = spacy.load(model_xx_name)
+
+if __name__ == "__main__":
+    text = """
+        This is a very long text. Actually, not that long but still made of a few sentences.
+    """
+    sentences = [sent.text.strip() for sent in nlp(text).sents if sent.text.strip()]
+    print(f"Nb of sentences: {len(sentences)}")
+    for i, sent in enumerate(sentences):
+        print(f"{i:2}: {sent}")

model_translation.py

@@ -0,0 +1,69 @@
+"""
+File: model_translation.py
+
+Description: 
+   Loading models for text translations (EN->FR, FR->EN)
+
+Author: Didier Guillevic
+Date: 2024-03-16
+"""
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+src_langs = set(["ar", "en", "fa", "fr", "he", "ja", "zh"])
+model_names = {
+    "ar": "Helsinki-NLP/opus-mt-ar-en",
+    "en": "Helsinki-NLP/opus-mt-en-fr",
+    "fa": "Helsinki-NLP/opus-mt-tc-big-fa-itc",
+    "fr": "Helsinki-NLP/opus-mt-fr-en",
+    "he": "Helsinki-NLP/opus-mt-tc-big-he-en",
+    "ja": "Helsinki-NLP/opus-mt-jap-en",
+    "zh": "Helsinki-NLP/opus-mt-zh-en",
+}
+
+# Registry for all loaded bilingual models
+tokenizer_model_registry = {}
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def get_tokenizer_model_for_src_lang(src_lang: str) -> (AutoTokenizer, AutoModelForSeq2SeqLM):
+    """
+    Return the (tokenizer, model) for a given source language.
+    """
+    src_lang = src_lang.lower()
+
+    # Already loaded?
+    if src_lang in tokenizer_model_registry:
+        return tokenizer_model_registry.get(src_lang)
+
+    # Load tokenizer and model
+    model_name = model_names.get(src_lang)
+    if not model_name:
+        raise Exception(f"No model defined for language: {src_lang}")
+    
+    # We will leave the models on the CPU (for now)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    if model.config.torch_dtype != torch.float16:
+        model = model.half()
+    model = model.to(device)
+    tokenizer_model_registry[src_lang] = (tokenizer, model)
+
+    return (tokenizer, model)
+
+# Max number of words for given input text
+# - Usually 512 tokens (max position encodings, as well as max length)
+# - Let's set to some number of words somewhat lower than that threshold
+# - e.g. 200 words
+max_words_per_chunk = 200
+
+#
+# Multilingual translation model
+#
+model_MADLAD_name = "google/madlad400-3b-mt"
+#model_MADLAD_name = "google/madlad400-7b-mt-bt"
+tokenizer_multilingual = AutoTokenizer.from_pretrained(model_MADLAD_name, use_fast=True)
+model_multilingual = AutoModelForSeq2SeqLM.from_pretrained(
+    model_MADLAD_name, device_map="auto", torch_dtype=torch.float16)
+