~ app.py

app.py

@@ -9,11 +9,11 @@ tgt_lang="zu"
 tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang=src_lang)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
 
-def translate(text):
+def translate(text, num_beams=4, num_return_sequences=4):
   inputs = tokenizer(text, return_tensors="pt")
 
   translated_tokens = model.generate(
-      **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=4, num_return_sequences=4
+      **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=num_beams, num_return_sequences=num_return_sequences
   )
 
   translations = []
@@ -23,7 +23,9 @@ def translate(text):
   return  "\n".join(["• " + translation for translation in translations])
 
 output = gr.outputs.Textbox()
-iface = gr.Interface(fn=translate, inputs="text", outputs=output)
+num_beams = gr.inputs.Slider(2, 64, 1, 2, label="Number of beams", default=4)
+num_return_sequences = gr.inputs.Slider(2, 64, 1, 2, label="Number of return sequences", default=4)
+iface = gr.Interface(fn=translate, inputs=["text", num_beams, num_return_sequences], outputs=output)
 iface.launch()
 
 # import gradio as gr