Update app.py

app.py

@@ -1,29 +1,12 @@
-import os
-import torch
-import gradio as gr
-import time
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from flores200_codes import flores_codes
-
 
-def load_models():
-    # build model and tokenizer
-    model_name_dict = {'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',
-                  }
-
-    model_dict = {}
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import torch
 
-    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
+# this model was loaded from https://hf.co/models
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-1.3B").to(device)
+tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-1.3B")
 
-    return model_dict
 
 LANGS = ["pes_Arab", "ckb_Arab", "eng_Latn"]
 langs_dict = {
@@ -36,13 +19,7 @@ def translate(text, src_lang, tgt_lang):
     """
     Translate the text from source lang to target lang
     """
-
-    if len(model_dict) == 2:
-        model_name = 'nllb-3.3B'
-    model = model_dict[model_name + '_model']
-    tokenizer = model_dict[model_name + '_tokenizer']
-
-    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, src_lang=langs_dict[src_lang], tgt_lang=langs_dict[tgt_lang], max_length=400, device="cpu")
+    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, src_lang=langs_dict[src_lang], tgt_lang=langs_dict[tgt_lang], max_length=400, device=device)
     result = translation_pipeline(text)
     return result[0]['translation_text']
 
@@ -76,12 +53,7 @@ def add_line(input_path, output_path):
     return output_path
 
 if __name__ == '__main__':
-    print('\tinit models')
-
-    #global model_dict
 
-    #model_dict = load_models()
-    
     interface = gr.Interface(
     fn=file_translate,
     inputs=[