Restructuring. Remove examples.

app.py

@@ -6,70 +6,17 @@ Description: Translate text...
 Author: Didier Guillevic
 Date: 2024-09-07
 """
-import spaces
-import gradio as gr
-import langdetect
-from typing import List
-
 import logging
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
+import gradio as gr
+import langdetect
+
 from deep_translator import GoogleTranslator
 from model_spacy import nlp_xx
 
-import model_translation as translation
-
-
-def detect_language(text):
-    lang = langdetect.detect(text)
-    return lang
-
-
-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)
-
-    return chunks
+import model_translation
 
 
 def translate_with_Helsinki(
@@ -106,82 +53,31 @@ def translate_with_Helsinki(
     return '\n'.join(translated_chunks)
 
 
-@spaces.GPU
-def translate_with_m2m100(
-        chunks: List[str],
-        src_lang: str,
-        tgt_lang: str) -> str:
-    """Translate with the m2m100 model
-    """
-    m2m100 = translation.ModelM2M100()
-    m2m100.tokenizer.src_lang = src_lang
-
-    translated_chunks = []
-    for chunk in chunks:
-        input_ids = m2m100.tokenizer(
-            chunk, return_tensors="pt").input_ids.to(m2m100.model.device)
-        outputs = m2m100.model.generate(
-            input_ids=input_ids,
-            forced_bos_token_id=m2m100.tokenizer.get_lang_id(tgt_lang))
-        translated_chunk = m2m100.tokenizer.batch_decode(
-            outputs, skip_special_tokens=True)[0]
-        translated_chunks.append(translated_chunk)
-    
-    return '\n'.join(translated_chunks)
-
-
-@spaces.GPU
-def translate_with_MADLAD(
-        chunks: List[str],
-        tgt_lang: str,
-        input_max_length: int=512,
-        output_max_length: int=512) -> str:
-    """Translate with Google MADLAD model
-    """
-    madlad = translation.ModelMADLAD()
-
-    translated_chunks = []
-    for chunk in chunks:
-        input_text = f"<2{tgt_lang}> {chunk}"
-        #logger.info(f" Translating: {input_text[:30]}")
-        input_ids = madlad.tokenizer(
-            input_text, return_tensors="pt", max_length=input_max_length,
-            truncation=True, padding="longest").input_ids.to(madlad.model.device)
-        outputs = madlad.model.generate(
-            input_ids=input_ids, max_length=output_max_length)
-        translated_chunk = madlad.tokenizer.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)  
+        src_lang: str,
+        tgt_lang: str
+    ) -> str:
+    """Translate the given text into English or French
     """
-    src_lang = src_lang if (src_lang and src_lang != "auto") else detect_language(text)
-    tgt_lang = 'en' # Default "easy" language
-
-    chunks = build_text_chunks(text, src_lang, sents_per_chunk)
-    chunks = [text, ]
+    # src_lang among the supported languages?
+    # - make sure src_lang is not None
+    src_lang = src_lang if (src_lang and src_lang != "auto") else langdetect.detect(text)
+    if src_lang not in model_translation.language_codes.values():
+        logging.error(f"Language detected {src_lang} not among supported language")
+    
+    # tgt_lang: make sure it is not None. Default to 'en' if not set.
+    if tgt_lang not in model_translation.tgt_language_codes.values():
+        tgt_lang = 'en'
 
-    #translated_text_Helsinki = translate_with_Helsinki(
-    #    chunks, src_lang, tgt_lang, input_max_length, output_max_length)
-    #translated_text_m2m100 = translate_with_m2m100(chunks, src_lang, tgt_lang)
-    translated_text_MADLAD = translate_with_MADLAD(chunks, tgt_lang)
+    # translate
+    m2m100 = model_translation.ModelM2M100()
+    translated_text_m2m100 = m2m100.translate(text, src_lang, tgt_lang)
     translated_text_google_translate = GoogleTranslator(
         source='auto', target='en').translate(text=text)
 
     return (
-        #translated_text_Helsinki,
-        #translated_text_m2m100,
-        translated_text_MADLAD,
+        translated_text_m2m100,
         translated_text_google_translate
     )
 
@@ -192,77 +88,54 @@ def translate_text(
 with gr.Blocks() as demo:
 
     gr.Markdown("""
-        ## Text translation v0.0.2 (small paragraph, multilingual)
+        ## Text translation v0.0.3
     """)
+    # Input
     input_text = gr.Textbox(
-        lines=15,
+        lines=6,
         placeholder="Enter text to translate",
         label="Text to translate",
-        render=False
+        render=True
     )
-    #output_text_Helsinki = gr.Textbox(
-    #    lines=6,
-    #    label="Bilingual translation model (Helsinki NLP)",
-    #    render=False
-    #)
-    #output_text_m2m100 = gr.Textbox(
-    #    lines=6,
-    #    label="Facebook m2m100 (1.2B)",
-    #    render=False
-    #)
-    output_text_MADLAD = gr.Textbox(
+
+    # Output
+    output_text_m2m100 = gr.Textbox(
         lines=6,
-        label="Google MADLAD400 (3B)",
-        render=False
+        label="Facebook m2m100 (1.2B)",
+        render=True
     )
     output_text_google_translate = gr.Textbox(
         lines=6,
         label="Google Translate",
-        render=False
+        render=True
     )
 
-    # 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", "zh"], value="auto",
-        label="Source language",
-        render=False
-    )
+    # Source and target languages
+    with gr.Row():
+        src_lang = gr.Radio(
+            choices=model_translation.language_codes.items(),
+            value="auto",
+            label="Source language",
+            render=True
+        )
+        tgt_lang = gr.Radio(
+            choices=model_translation.tgt_language_codes.items(),
+            value="en",
+            label="Target language",
+            render=True
+        )
 
-    # Examples
-    examples = [
-        ["ریچارد مور، رئیس سازمان مخفی اطلاعاتی بریتانیا (ام‌آی‌۶) در دیدار ویلیام برنز، رئیس سازمان اطلاعات مرکزی آمریکا (سیا) گفت همچنان احتمال اقدام ایران علیه اسرائیل در واکنش به ترور اسماعیل هنیه، رهبر حماس وجود دارد. آقای برنز نیز در این دیدار فاش کرد که در سال اول جنگ اوکراین، «خطر واقعی» وجود داشت که روسیه به استفاده از «تسلیحات هسته‌ای تاکتیکی» متوسل شود. این دو مقام امنیتی هشدار دادند که «نظم جهانی» از زمان جنگ سرد تا کنون تا این حد «در معرض تهدید» نبوده است.", "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"],
-        ["張先生稱，奇瑞已經凖備在西班牙生產汽車，並決心採取「本地化」的方式進入歐洲市場。此外，他也否認該公司的出口受益於不公平補貼。奇瑞成立於1997年，是中國最大的汽車公司之一。它已經是中國最大的汽車出口商，並且制定了進一步擴張的野心勃勃的計劃。", "zh"],
-        #["ברוכה הבאה, קיטי: בית הקפה החדש בלוס אנג'לס החתולה האהובה והחברים שלה מקבלים בית קפה משלהם בשדרות יוניברסל סיטי, שם תוכלו למצוא מגוון של פינוקים מתוקים – החל ממשקאות ועד עוגות", "he"],
-    ]
-    
-    gr.Interface(
+    # Submit button
+    translate_btn = gr.Button("Translate")
+    translate_btn.click(
         fn=translate_text,
-        inputs=[input_text, src_lang],
-        outputs=[
-            #output_text_Helsinki,
-            #output_text_m2m100,
-            output_text_MADLAD,
-            output_text_google_translate,
-        ],
-        additional_inputs=[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=False
+        inputs=[input_text, src_lang, tgt_lang],
+        outputs=[output_text_m2m100, output_text_google_translate]
     )
 
     with gr.Accordion("Documentation", open=False):
         gr.Markdown("""
-            - Models: serving Facebook M2M100 and Google MADLAD models.
-            - Basic: processing of long paragraph / document to be enhanced.
-            - Most examples are copy/pasted from BBC news international web sites.
+            - Models: serving Facebook M2M100 and Google Translate.
         """)
 
 if __name__ == "__main__":

model_translation.py

@@ -7,17 +7,188 @@ Description:
 Author: Didier Guillevic
 Date: 2024-03-16
 """
+import spaces
+
+import logging
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
 from transformers import BitsAndBytesConfig
 
+from model_spacy import nlp_xx as model_spacy
+
 quantization_config = BitsAndBytesConfig(
     load_in_8bit=True,
     llm_int8_threshold=200.0 # https://discuss.huggingface.co/t/correct-usage-of-bitsandbytesconfig/33809/5
 )
 
+# The 100 languages supported by the facebook/m2m100_418M model
+# https://huggingface.co/facebook/m2m100_418M
+# plus the 'AUTOMATIC' option where we will use a language detector.
+language_codes = {
+    'AUTOMATIC': 'auto',
+    'Afrikaans (af)': 'af',
+    'Albanian (sq)': 'sq',
+    'Amharic (am)': 'am',
+    'Arabic (ar)': 'ar',
+    'Armenian (hy)': 'hy',
+    'Asturian (ast)': 'ast',
+    'Azerbaijani (az)': 'az',
+    'Bashkir (ba)': 'ba',
+    'Belarusian (be)': 'be',
+    'Bengali (bn)': 'bn',
+    'Bosnian (bs)': 'bs',
+    'Breton (br)': 'br',
+    'Bulgarian (bg)': 'bg',
+    'Burmese (my)': 'my',
+    'Catalan; Valencian (ca)': 'ca',
+    'Cebuano (ceb)': 'ceb',
+    'Central Khmer (km)': 'km',
+    'Chinese (zh)': 'zh',
+    'Croatian (hr)': 'hr',
+    'Czech (cs)': 'cs',
+    'Danish (da)': 'da',
+    'Dutch; Flemish (nl)': 'nl',
+    'English (en)': 'en',
+    'Estonian (et)': 'et',
+    'Finnish (fi)': 'fi',
+    'French (fr)': 'fr',
+    'Fulah (ff)': 'ff',
+    'Gaelic; Scottish Gaelic (gd)': 'gd',
+    'Galician (gl)': 'gl',
+    'Ganda (lg)': 'lg',
+    'Georgian (ka)': 'ka',
+    'German (de)': 'de',
+    'Greeek (el)': 'el',
+    'Gujarati (gu)': 'gu',
+    'Haitian; Haitian Creole (ht)': 'ht',
+    'Hausa (ha)': 'ha',
+    'Hebrew (he)': 'he',
+    'Hindi (hi)': 'hi',
+    'Hungarian (hu)': 'hu',
+    'Icelandic (is)': 'is',
+    'Igbo (ig)': 'ig',
+    'Iloko (ilo)': 'ilo',
+    'Indonesian (id)': 'id',
+    'Irish (ga)': 'ga',
+    'Italian (it)': 'it',
+    'Japanese (ja)': 'ja',
+    'Javanese (jv)': 'jv',
+    'Kannada (kn)': 'kn',
+    'Kazakh (kk)': 'kk',
+    'Korean (ko)': 'ko',
+    'Lao (lo)': 'lo',
+    'Latvian (lv)': 'lv',
+    'Lingala (ln)': 'ln',
+    'Lithuanian (lt)': 'lt',
+    'Luxembourgish; Letzeburgesch (lb)': 'lb',
+    'Macedonian (mk)': 'mk',
+    'Malagasy (mg)': 'mg',
+    'Malay (ms)': 'ms',
+    'Malayalam (ml)': 'ml',
+    'Marathi (mr)': 'mr',
+    'Mongolian (mn)': 'mn',
+    'Nepali (ne)': 'ne',
+    'Northern Sotho (ns)': 'ns',
+    'Norwegian (no)': 'no',
+    'Occitan (post 1500) (oc)': 'oc',
+    'Oriya (or)': 'or',
+    'Panjabi; Punjabi (pa)': 'pa',
+    'Persian (fa)': 'fa',
+    'Polish (pl)': 'pl',
+    'Portuguese (pt)': 'pt',
+    'Pushto; Pashto (ps)': 'ps',
+    'Romanian; Moldavian; Moldovan (ro)': 'ro',
+    'Russian (ru)': 'ru',
+    'Serbian (sr)': 'sr',
+    'Sindhi (sd)': 'sd',
+    'Sinhala; Sinhalese (si)': 'si',
+    'Slovak (sk)': 'sk',
+    'Slovenian (sl)': 'sl',
+    'Somali (so)': 'so',
+    'Spanish (es)': 'es',
+    'Sundanese (su)': 'su',
+    'Swahili (sw)': 'sw',
+    'Swati (ss)': 'ss',
+    'Swedish (sv)': 'sv',
+    'Tagalog (tl)': 'tl',
+    'Tamil (ta)': 'ta',
+    'Thai (th)': 'th',
+    'Tswana (tn)': 'tn',
+    'Turkish (tr)': 'tr',
+    'Ukrainian (uk)': 'uk',
+    'Urdu (ur)': 'ur',
+    'Uzbek (uz)': 'uz',
+    'Vietnamese (vi)': 'vi',
+    'Welsh (cy)': 'cy',
+    'Western Frisian (fy)': 'fy',
+    'Wolof (wo)': 'wo',
+    'Xhosa (xh)': 'xh',
+    'Yiddish (yi)': 'yi',
+    'Yoruba (yo)': 'yo',
+    'Zulu (zu)': 'zu'
+}
+
+tgt_language_codes = {
+    'English (en)': 'en',
+    'French (fr)': 'fr'
+}
+
+
+def build_text_chunks(
+        text: str,
+        sents_per_chunk: int=5,
+        words_per_chunk=200) -> list[str]:
+    """Split a given text into chunks with at most sents_per_chnks and words_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 words_per_chunk
+    """
+    # Split text into sentences...
+    sentences = [
+        sent.text.strip() for sent in model_spacy(text).sents if sent.text.strip()
+    ]
+    logger.info(f"TEXT: {text[:25]}, 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 > 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)
+
+    return chunks
+
+
 class Singleton(type):
     _instances = {}
     def __call__(cls, *args, **kwargs):
@@ -25,8 +196,11 @@ class Singleton(type):
             cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
         return cls._instances[cls]
 
+
 class ModelM2M100(metaclass=Singleton):
     """Loads an instance of the M2M100 model.
+
+    Model: https://huggingface.co/facebook/m2m100_1.2B
     """
     def __init__(self):
         self._model_name = "facebook/m2m100_1.2B"
@@ -35,10 +209,47 @@ class ModelM2M100(metaclass=Singleton):
             self._model_name,
             device_map="auto",
             torch_dtype=torch.float16,
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
+            quantization_config=quantization_config
         )
         self._model = torch.compile(self._model)
     
+    @spaces.GPU
+    def translate(
+            self,
+            text: str,
+            src_lang: str,
+            tgt_lang: str,
+            chunk_text: bool=True,
+            sents_per_chunk: int=5,
+            words_per_chunk: int=200
+        ) -> str:
+        """Translate the given text from src_lang to tgt_lang.
+
+        The text will be split into chunks to ensure the chunks fit into the 
+        model input_max_length (usually 512 tokens).
+        """
+        chunks = [text,]
+        if chunk_text:
+            chunks = build_text_chunks(text, sents_per_chunk, words_per_chunk)
+        
+        self._tokenizer.src_lang = src_lang
+
+        translated_chunks = []
+        for chunk in chunks:
+            input_ids = self._tokenizer(
+                chunk,
+                return_tensors="pt").input_ids.to(self._model.device)
+            outputs = self._model.generate(
+                input_ids=input_ids,
+                forced_bos_token_id=self._tokenizer.get_lang_id(tgt_lang))
+            translated_chunk = self._tokenizer.batch_decode(
+                outputs,
+                skip_special_tokens=True)[0]
+            translated_chunks.append(translated_chunk)
+
+        return '\n'.join(translated_chunks)
+
     @property
     def model_name(self):
         return self._model_name
@@ -51,11 +262,20 @@ class ModelM2M100(metaclass=Singleton):
     def model(self):
         return self._model
 
+    @property
+    def device(self):
+        return self._model.device
+
+
 class ModelMADLAD(metaclass=Singleton):
     """Loads an instance of the Google MADLAD model (3B).
+
+    Model: https://huggingface.co/google/madlad400-3b-mt
     """
     def __init__(self):
         self._model_name = "google/madlad400-3b-mt"
+        self._input_max_length = 512 # config.json n_positions
+        self._output_max_length = 512 # config.json n_positions
         self._tokenizer = AutoTokenizer.from_pretrained(
             self.model_name, use_fast=True
         )
@@ -68,6 +288,44 @@ class ModelMADLAD(metaclass=Singleton):
         )
         self._model = torch.compile(self._model)
     
+    @spaces.GPU
+    def translate(
+            self,
+            text: str,
+            tgt_lang: str,
+            chunk_text: True,
+            sents_per_chunk: int=5,
+            words_per_chunk: int=5
+        ) -> str:
+        """Translate given text into the target language.
+
+        The text will be split into chunks to ensure the chunks fit into the 
+        model input_max_length (usually 512 tokens).
+        """
+        chunks = [text,]
+        if chunk_text:
+            chunks = build_text_chunks(text, sents_per_chunk, words_per_chunk)
+        
+        translated_chunks = []
+        for chunk in chunks:
+            input_text = f"<2{tgt_lang}> {chunk}"
+            logger.info(f" Translating: {input_text[:50]}")
+            input_ids = self._tokenizer(
+                input_text,
+                return_tensors="pt",
+                max_length=self._input_max_length,
+                truncation=True,
+                padding="longest").input_ids.to(self._model.device)
+            outputs = self._model.generate(
+                input_ids=input_ids,
+                max_length=self._output_max_length)
+            translated_chunk = self._tokenizer.decode(
+                outputs[0],
+                skip_special_tokens=True)
+            translated_chunks.append(translated_chunk)
+    
+        return '\n'.join(translated_chunks)
+
     @property
     def model_name(self):
         return self._model_name
@@ -79,6 +337,10 @@ class ModelMADLAD(metaclass=Singleton):
     @property
     def model(self):
         return self._model
+    
+    @property
+    def device(self):
+        return self._model.device
 
 
 # Bi-lingual individual models