Update app.py

app.py

@@ -1,11 +1,7 @@
-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',
@@ -25,8 +21,7 @@ def load_models():
 
     return model_dict
 
-
-def translation(source, target, text):
+def translation(source, target, text, model_name="nllb-distilled-600M"):
     if len(model_dict) == 2:
         model_name = 'nllb-distilled-600M'
 
@@ -51,35 +46,21 @@ def translation(source, target, text):
 
 
 if __name__ == '__main__':
-    print('\tinit models')
-
-    global model_dict
-
     model_dict = load_models()
-    
-    # define gradio demo
+
     lang_codes = list(flores_codes.keys())
-    #inputs = [gr.Image.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
-    Image = [gr.Image.Dropdown(lang_codes, default='English', label='Source'),
-              gr.Image.Dropdown(lang_codes, default='Korean', label='Target'),
-              gr.Image.Textbox(lines=5, label="Input text"),
-              ]
 
-    Label = gr.Label.JSON()
+    inputs = [
+        gr.components.Dropdown(lang_codes, default='English', label='Source'),
+        gr.components.Dropdown(lang_codes, default='Korean', label='Target'),
+        gr.components.Textbox(lines=5, label="Input text"),
+        gr.components.Dropdown(["nllb-distilled-600M"], default="nllb-distilled-600M", label="Model"),
+    ]
 
-    title = "NLLB distilled 600M demo"
+    outputs = gr.components.JSON()
 
+    title = "NLLB distilled 600M demo"
     demo_status = "Demo is running on CPU"
     description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"
-    examples = [
-    ['English', 'Korean', 'Hi. nice to meet you']
-    ]
-
-    gr.Interface(translation,
-                 Image,
-                 Label,
-                 title=title,
-                 description=description,
-                 ).launch()
-
 
+    gr.Interface(translation, inputs, outputs, title=title, description=description).launch()