add punctuation normalization and load the tokenizer only once

app.py

@@ -1,5 +1,6 @@
 import spaces
 import gradio as gr
+from sacremoses import MosesPunctNormalizer
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from flores import code_mapping
 import platform
@@ -28,12 +29,11 @@ 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")
 
 # cache function
 @lru_cache(maxsize=100)
@@ -44,10 +44,17 @@ def translate(text: str, src_lang: str, tgt_lang: str):
         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 = []
@@ -66,7 +73,7 @@ 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,

requirements.txt

@@ -3,4 +3,5 @@ transformers
 torch
 gradio==4.32.2
 spaces
-nltk
+nltk
+sacremoses