[update]add main

main.py

@@ -13,27 +13,32 @@ import gradio as gr
 from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 
 
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--pretrained_model_name_or_path",
-        default="facebook/m2m100_418M",
-        type=str
-    )
-    args = parser.parse_args()
-    return args
-
-
 def main():
-    args = get_args()
-
-    model = M2M100ForConditionalGeneration.from_pretrained(args.pretrained_model_name_or_path)
-    tokenizer = M2M100Tokenizer.from_pretrained(args.pretrained_model_name_or_path)
+    model_dict = {
+        "facebook/m2m100_418M": {
+            "model": M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M"),
+            "tokenizer": M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
+        }
+    }
 
     def multilingual_translate(src_text: str,
                                src_lang: str,
                                tgt_lang: str,
+                               model_name: str,
                                ):
+        model_group = model_dict.get(model_name)
+        if model_group is None:
+            for k, mg in model_dict.items():
+                del mg["model"]
+            model_dict[model_name] = {
+                "model": M2M100ForConditionalGeneration.from_pretrained(model_name),
+                "tokenizer": M2M100Tokenizer.from_pretrained(model_name)
+            }
+            model_group = model_dict[model_name]
+
+        model = model_group["model"]
+        tokenizer = model_group["tokenizer"]
+
         tokenizer.src_lang = src_lang
         encoded_src = tokenizer(src_text, return_tensors="pt")
         generated_tokens = model.generate(**encoded_src,
@@ -77,13 +82,19 @@ Yiddish (yi), Yoruba (yo), Chinese (zh), Zulu (zu)
             "Hello world!",
             "en",
             "zh",
+            "facebook/m2m100_418M",
         ],
     ]
 
+    model_choices = [
+        "facebook/m2m100_418M",
+        "facebook/m2m100_1.2B"
+    ]
     inputs = [
         gr.Textbox(lines=4, value="", label="Input Text"),
         gr.Textbox(lines=1, value="", label="Source Language"),
         gr.Textbox(lines=1, value="", label="Target Language"),
+        gr.Dropdown(choices=model_choices, label="model_name")
     ]
 
     output = gr.Textbox(lines=4, label="Output Text")