Update app.py

app.py

@@ -24,7 +24,7 @@ models = ["Helsinki-NLP", "QUICKMT", "Argos", "Google", "HPLT", "HPLT-OPUS",
           "Helsinki-NLP/opus-mt-tc-bible-big-roa-deu_eng_fra_por_spa", "Helsinki-NLP/opus-mt-tc-bible-big-deu_eng_fra_por_spa-roa", "Helsinki-NLP/opus-mt-tc-bible-big-roa-en",
           "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B", "facebook/nllb-200-1.3B", "facebook/nllb-200-3.3B",
           "facebook/mbart-large-50-many-to-many-mmt", "facebook/mbart-large-50-one-to-many-mmt", "facebook/mbart-large-50-many-to-one-mmt",
-          "facebook/m2m100_418M", "facebook/m2m100_1.2B", "Lego-MT/Lego-MT",
+          "facebook/m2m100_418M", "facebook/m2m100_1.2B", "alirezamsh/small100", "Lego-MT/Lego-MT",
           "bigscience/mt0-small", "bigscience/mt0-base", "bigscience/mt0-large", "bigscience/mt0-xl",
           "bigscience/bloomz-560m", "bigscience/bloomz-1b1", "bigscience/bloomz-1b7", "bigscience/bloomz-3b",
           "google-t5/t5-small", "google-t5/t5-base", "google-t5/t5-large",
@@ -102,6 +102,16 @@ class Translators:
         outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         outputs = ''.join(outputs) if isinstance(outputs, list) else outputs
         return outputs
+
+    def smallonehundred(self):
+        from transformers import M2M100ForConditionalGeneration
+        from tokenization_small100 import SMALL100Tokenizer
+        model = M2M100ForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = SMALL100Tokenizer.from_pretrained(self.model_name)
+        tokenizer.tgt_lang = self.tl
+        encoded_sl = tokenizer(self.input_text, return_tensors="pt")
+        generated_tokens = model.generate(**encoded_sl, max_length=256, num_beams=5)
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)0]
     
     def hplt(self, opus = False):
         # langs = ['ar', 'bs', 'ca', 'en', 'et', 'eu', 'fi', 'ga', 'gl', 'hi', 'hr', 'is', 'mt', 'nn', 'sq', 'sw', 'zh_hant']
@@ -109,6 +119,9 @@ class Translators:
                   'en-ga', 'en-gl', 'en-hi', 'en-hr', 'en-is', 'en-mt', 'en-nn', 'en-sq', 'en-sw',
                   'en-zh_hant', 'et-en', 'eu-en', 'fi-en', 'ga-en', 'gl-en', 'hi-en', 'hr-en',
                   'is-en', 'mt-en', 'nn-en', 'sq-en', 'sw-en', 'zh_hant-en']
+        lang_map = {"zh": "zh_hant"}
+        self.sl = lang_map.get(self.sl, self.sl)
+        self.tl = lang_map.get(self.tl, self.tl)
         if opus:
             hplt_model = f'HPLT/translate-{self.sl}-{self.tl}-v1.0-hplt_opus' # HPLT/translate-en-hr-v1.0-hplt_opus
         else:
@@ -534,6 +547,9 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
         
         elif "m2m" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).mtom()
+        
+        elif "small100" in model_name.lower():
+            translated_text = Translators(model_name, sl, tl, input_text).smallonehundred()
 
         elif "lego" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).LegoMT()