Optimize the preprocessing and generation

app.py

@@ -1,5 +1,7 @@
 import spaces
 import gradio as gr
+from sacremoses import MosesPunctNormalizer
+from stopes.pipelines.monolingual.utils.sentence_split import get_split_algo
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from flores import code_mapping
 import platform
@@ -28,28 +30,47 @@ def load_model():
 model = load_model()
 
 
-def load_tokenizer(src_lang, tgt_lang):
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME, src_lang=code_mapping[src_lang], tgt_lang=code_mapping[tgt_lang]
-    )
-    return tokenizer
+# Loading the tokenizer once, because re-loading it takes about 1.5 seconds each time
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+
+punct_normalizer = MosesPunctNormalizer(lang="en")
+
+
+@lru_cache(maxsize=202)
+def get_language_specific_sentence_splitter(language_code):
+    short_code = language_code[:3]
+    splitter = get_split_algo(short_code, "default")
+    return splitter
 
 
 # cache function
 @lru_cache(maxsize=100)
 def translate(text: str, src_lang: str, tgt_lang: str):
-    return _translate(text, src_lang,tgt_lang )
+    if not src_lang:
+        raise gr.Error("The source language is empty! Please choose it in the dropdown list.")
+    if not tgt_lang:
+        raise gr.Error("The target language is empty! Please choose it in the dropdown list.")
+    return _translate(text, src_lang, tgt_lang)
+
 
 # Only assign GPU if cache not used
 @spaces.GPU
 def _translate(text: str, src_lang: str, tgt_lang: str):
-    tokenizer = load_tokenizer(src_lang, tgt_lang)
+    src_code = code_mapping[src_lang]
+    tgt_code = code_mapping[tgt_lang]
+    tokenizer.src_lang = src_code
+    tokenizer.tgt_lang = tgt_code
+
+    # normalizing the punctuation first
+    text = punct_normalizer.normalize(text)
 
     paragraphs = text.split("\n")
     translated_paragraphs = []
 
     for paragraph in paragraphs:
-        sentences = nltk.sent_tokenize(paragraph)
+        splitter = get_language_specific_sentence_splitter(src_code)
+        sentences = list(splitter(paragraph))
         translated_sentences = []
 
         for sentence in sentences:
@@ -62,9 +83,12 @@ def _translate(text: str, src_lang: str, tgt_lang: str):
             )
             translated_chunk = model.generate(
                 input_ids=torch.tensor([input_tokens]).to(device),
-                forced_bos_token_id=tokenizer.convert_tokens_to_ids(code_mapping[tgt_lang]),
+                forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_code),
                 max_length=len(input_tokens) + 50,
                 num_return_sequences=1,
+                num_beams=5,
+                no_repeat_ngram_size=4,  # repetition blocking works better if this number is below num_beams  
+                renormalize_logits=True,  # recompute token probabilities after banning the repetitions
             )
             translated_chunk = tokenizer.decode(
                 translated_chunk[0], skip_special_tokens=True

flores.py

@@ -10,7 +10,7 @@ code_mapping = {
     "Amharic": "amh_Ethi",
     "North Levantine Arabic": "apc_Arab",
     "Modern Standard Arabic": "arb_Arab",
-    "Modern Standard Arabic (Romanized)": "arb_Latn",
+    # "Modern Standard Arabic (Romanized)": "arb_Latn",  # it is in FLORES, but not in NLLB
     "Najdi Arabic": "ars_Arab",
     "Moroccan Arabic": "ary_Arab",
     "Egyptian Arabic": "arz_Arab",
@@ -115,7 +115,7 @@ code_mapping = {
     "Maithili": "mai_Deva",
     "Malayalam": "mal_Mlym",
     "Marathi": "mar_Deva",
-    "Minangkabau (Arabic script)": "min_Arab",
+    # "Minangkabau (Arabic script)": "min_Arab",   # it is in FLORES, but not in NLLB
     "Minangkabau (Latin script)": "min_Latn",
     "Macedonian": "mkd_Cyrl",
     "Plateau Malagasy": "plt_Latn",
@@ -149,7 +149,7 @@ code_mapping = {
     "Russian": "rus_Cyrl",
     "Sango": "sag_Latn",
     "Sanskrit": "san_Deva",
-    "Santali": "sat_Olck",
+    "Santali": "sat_Beng",  # It is called sat_Olck in FLORES, but (less correctly sat_Beng in NLLB)
     "Sicilian": "scn_Latn",
     "Shan": "shn_Mymr",
     "Sinhala": "sin_Sinh",

requirements.txt

@@ -3,4 +3,6 @@ transformers
 torch
 gradio==4.32.2
 spaces
-nltk
+nltk
+sacremoses
+stopes[mono] @ git+https://github.com/facebookresearch/stopes@better-sentence-splitters