Update app.py

app.py

@@ -4,32 +4,38 @@ from multilingual_translation import translate
 from utils import lang_ids
 import gradio as gr
 
-model_list = [
+biogpt_model_list = [
     "microsoft/biogpt",
     "microsoft/BioGPT-Large-PubMedQA"
 ]
 
+lang_model_list = [
+    "facebook/m2m100_1.2B",
+    "facebook/m2m100_418M"
+]
+
 lang_list = list(lang_ids.keys())
 
-def translate_to_english(text, base_lang):
+def translate_to_english(text, lang_model_id, base_lang):
     if base_lang == "English":
         return text
     else:
         base_lang = lang_ids[base_lang]
-        new_text = translate("facebook/m2m100_418M", text, base_lang, "en")
+        new_text = translate(lang_model_id, text, base_lang, "en")
         return new_text
 
 def biogpt(
     prompt: str,
-    model_id: str,
+    biogpt_model_id: str,
+    lang_model_id:str,
     max_length: int = 25,
     num_return_sequences: int = 5,
     base_lang: str = "English"
 ):
     
     en_prompt = translate_to_english(prompt, base_lang)[0]
-    model = BioGptForCausalLM.from_pretrained(model_id)
-    tokenizer = BioGptTokenizer.from_pretrained(model_id)
+    model = BioGptForCausalLM.from_pretrained(biogpt_model_id)
+    tokenizer = BioGptTokenizer.from_pretrained(biogpt_model_id)
     generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
     set_seed(42)
     output = generator(en_prompt, max_length=max_length, num_return_sequences=num_return_sequences, do_sample=True)
@@ -50,7 +56,8 @@ inputs = [
     gr.Dropdown(model_list, value="microsoft/biogpt", label="Model ID"),
     gr.inputs.Slider(5, 100, 25, default=25, label="Max Length"),
     gr.inputs.Slider(1, 10, 5, default=5, label="Num Return Sequences"),
-    gr.Dropdown(lang_list, value="English", label="Base Language")
+    gr.Dropdown(lang_list, value="English", label="Base Language"),
+    gr.Dropdown(lang_model_list, value="microsoft/biogpt", label="Model ID"),
     ]
 
 outputs = [
@@ -59,8 +66,8 @@ outputs = [
 ]
 
 examples = [
-    ["COVID-19 is", "microsoft/biogpt", 25, 5, "English"],
-    ["Kanser", "microsoft/biogpt", 25, 5, "Turkish"]
+    ["COVID-19 is", "microsoft/biogpt", 25, 5, "English", "microsoft/biogpt"],
+    ["Kanser", "microsoft/biogpt", 25, 5, "Turkish", "facebook/m2m100_1.2B"]
 ]
 
 title = " BioGPT: Generative Pre-trained Transformer for Biomedical Text Generation and Mining"