Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import time
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from flores200_codes import flores_codes
+import vinai_translator
 
 def load_models():
     # build model and tokenizer
@@ -24,25 +25,38 @@ def load_models():
 
     return model_dict
 
-
 def translation(source, target, text):
-    if len(model_dict) == 2:
-        #model_name = 'nllb-distilled-600M'
-        model_name = 'nllb-distilled-1.3B'
-
-    start_time = time.time()
     source = flores_codes[source]
     target = flores_codes[target]
-
-    model = model_dict[model_name + '_model']
-    tokenizer = model_dict[model_name + '_tokenizer']
-
-    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
-    chunks = text.splitlines(True)
     output = ""
-    for chunk in chunks:
-        stchunk = translator(chunk+"<mask>", max_length=400, num_beams=5)
-        output += stchunk[0]['translation_text']+"\n"
+    
+    start_time = time.time()
+
+    if source = "vie_Latn" or target ="eng_Latn":
+        chunks = text.splitlines(True)
+        for chunk in chunks:
+            stchunk = vinai_translator.translate_vi2en(chunk)
+            output += stchunk+"\n"
+            pass
+    elif source = "eng_Latn" or target ="vie_Latn":
+        chunks = text.splitlines(True)
+        for chunk in chunks:
+            stchunk = vinai_translator.translate_en2vi(chunk)
+            output += stchunk+"\n"
+            pass
+    else:
+        if len(model_dict) == 2:
+            #model_name = 'nllb-distilled-600M'
+            model_name = 'nllb-distilled-1.3B'
+
+            model = model_dict[model_name + '_model']
+            tokenizer = model_dict[model_name + '_tokenizer']
+
+            translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
+            chunks = text.splitlines(True)
+            for chunk in chunks:
+                stchunk = translator(chunk+"<mask>", max_length=400, num_beams=5)
+                output += stchunk[0]['translation_text']+"\n"
 
     end_time = time.time()
 
@@ -52,7 +66,6 @@ def translation(source, target, text):
               'result': output}
     return result
 
-
 if __name__ == '__main__':
     print('\tinit models')