Update app.py

app.py

@@ -27,22 +27,21 @@ tokens = [tokenizer.decode(i) for i in range(460)]
 lang_codes = [token[2:-1] for token in tokens if token.startswith("<2")]
 
 
-def translate(input_text, target_language):
+def translate(input_text, target_language, beam_size, no_repeat_ngram_size, repetition_penalty):
     input_tokens = tokenizer.encode(f"<2{target_language}> {input_text}", out_type=str)
     results = translator.translate_batch(
         [input_tokens],
         batch_type="tokens",
-        # max_batch_size=1024,
-        beam_size=1,
-        no_repeat_ngram_size=1,
-        # repetition_penalty=2,
+        beam_size=beam_size,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        repetition_penalty=repetition_penalty,
     )
     translated_sentence = tokenizer.decode(results[0].hypotheses[0])
     return translated_sentence
 
 
-def translate_interface(input_text, target_language):
-    translated_text = translate(input_text, target_language)
+def translate_interface(input_text, target_language, beam_size, no_repeat_ngram_size, repetition_penalty):
+    translated_text = translate(input_text, target_language, beam_size, no_repeat_ngram_size, repetition_penalty)
     return translated_text
 
 
@@ -51,12 +50,15 @@ input_text = gr.Textbox(
     value="Imagine a world in which every single person on the planet is given free access to the sum of all human knowledge.",
 )
 target_language = gr.Dropdown(lang_codes, value="ml", label="Target Language")
+beam_size = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="Beam Size")
+no_repeat_ngram_size = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="No Repeat N-gram Size")
+repetition_penalty = gr.Slider(minimum=1.0, maximum=3.0, value=1.0, step=0.1, label="Repetition Penalty")
 output_text = gr.Textbox(label="Translated Text")
 
 gr.Interface(
     title=title,
     description=description,
     fn=translate_interface,
-    inputs=[input_text, target_language],
+    inputs=[input_text, target_language, beam_size, no_repeat_ngram_size, repetition_penalty],
     outputs=output_text,
-).launch()
+).launch()