Update the `app.py`

app.py

@@ -8,13 +8,38 @@ from flores200_codes import flores_codes
 
 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',
+    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 = {}
+
+    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
+
+
+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',
-                  # 'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M', 
                   }
 
     model_dict = {}
@@ -31,7 +56,7 @@ def load_models():
 
 def translation(source, target, text):
     if len(model_dict) == 2:
-        model_name = 'nllb-distilled-1.3B'
+        model_name = 'nllb-distilled-600M'
 
     start_time = time.time()
     source = flores_codes[source]
@@ -53,6 +78,7 @@ def translation(source, target, text):
     return result
 
 
+
 if __name__ == '__main__':
     print('\tinit models')
 
@@ -70,7 +96,7 @@ if __name__ == '__main__':
 
     outputs = gr.outputs.JSON()
 
-    title = "NLLB distilled 600M demo endpoint"
+    title = "NLLB distilled 600M demo"
 
     demo_status = "Demo is running on CPU"
     description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"