Update app.py

app.py

@@ -28,23 +28,50 @@ pipe = pipeline(
 )
 
 
-def load_translation_model():
-    model_name = 'facebook/nllb-200-distilled-1.3B'
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    return model, tokenizer
-
-translation_model, translation_tokenizer = load_translation_model()
-
-
-def translate_text(text, source_language, target_language):
-    source_code = flores_codes[source_language]
-    target_code = flores_codes[target_language]
-
-    translator = pipeline('translation', model=translation_model, tokenizer=translation_tokenizer, src_lang=source_code, tgt_lang=target_code)
+def load_models():
+    # build model and tokenizer
+    model_name_dict = {
+                'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
+        #'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
+                  #'nllb-1.3B': 'facebook/nllb-200-1.3B',
+                  #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
+                  #'nllb-3.3B': 'facebook/nllb-200-3.3B',
+                  # 'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M', 
+                  }
+
+    model_dict = {}
+
+    for call_name, real_name in model_name_dict.items():
+        print('\tLoading model: %s' % call_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
+        tokenizer = AutoTokenizer.from_pretrained(real_name)
+        model_dict[call_name+'_model'] = model
+        model_dict[call_name+'_tokenizer'] = tokenizer
+
+    return model_dict
+
+def translation(source, target, text):
+    if len(model_dict) == 2:
+        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)
     output = translator(text, max_length=400)
-    return output[0]['translation_text']
 
+    end_time = time.time()
+
+    output = output[0]['translation_text']
+    result = {'inference_time': end_time - start_time,
+              'source': source,
+              'target': target,
+              'result': output}
+    return result
 
 
 def transcribe(inputs, task):
@@ -52,7 +79,7 @@ def transcribe(inputs, task):
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
 
     text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
-    translated_text = translation(source_lang, target_lang, text)["result"]
+    translated_text = translation(source, target, text)["result"]
     return  text, translated_text
 
 
@@ -109,16 +136,16 @@ def yt_transcribe(yt_url, task, max_filesize=75.0):
     inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
 
     text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
-    translated_text = translation(source_lang, target_lang, text)["result"]
+    translated_text = translation(source, target, text)["result"]
 
     return html_embed_str, text, translated_text
 
 
-lang_codes = list(flores_codes.keys())
-
 
 demo = gr.Blocks()
 
+lang_codes = list(flores_codes.keys())
+
 mf_transcribe = gr.Interface(
     fn=transcribe,
     inputs=[