return sequences added

app.py

@@ -6,13 +6,12 @@ import torch
 model = AutoModelForSeq2SeqLM.from_pretrained("Jayyydyyy/m2m100_418m_tokipona")
 tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
-LANGS = ["English", "toki pona"]
 LANG_CODES = {
     "English":"en",
     "toki pona":"tl"
 }
 
-def translate(text, src_lang, tgt_lang):
+def translate(text, src_lang, tgt_lang, candidates:int):
     """
     Translate the text from source lang to target lang
     """
@@ -30,8 +29,8 @@ def translate(text, src_lang, tgt_lang):
             'output_scores': True,
             'output_hidden_states': True,
             'length_penalty': 0.0,  # don't encourage longer or shorter output,
-            'num_return_sequences': 3,
-            'num_beams':3,
+            'num_return_sequences': candidates,
+            'num_beams':candidates,
             'forced_bos_token_id': tokenizer.lang_code_to_id[tgt]
         }
     
@@ -45,11 +44,12 @@ app = gr.Interface(
     fn=translate,
     inputs=[
         gr.components.Textbox(label="Text"),
-        gr.components.Dropdown(label="Source Language", choices=LANGS),
-        gr.components.Dropdown(label="Target Language", choices=LANGS),
+        gr.components.Dropdown(label="Source Language", choices=LANG_CODES.keys()),
+        gr.components.Dropdown(label="Target Language", choices=LANG_CODES.keys()),
+        gr.Slider(label="Number of return sequences", value=1, minimum=1, maximum=12)
     ],
     outputs=["text"],
-    examples=[["This is an example!", "English", "toki pona"]],
+    examples=[["This is an example statement. It will be translated from English to toki pona.", "English", "toki pona"]],
     cache_examples=False,
     title="A simple English / toki pona Neural Translation App",
     description="A simple English / toki pona Neural Translation App"