Update app.py

app.py

@@ -392,14 +392,14 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
     return original_tree
 
 @spaces.GPU
-def get_beam_search_html(input_text, number_steps, number_beams, length_penalty):
+def get_beam_search_html(input_text, number_steps, number_beams, length_penalty, number_sequences):
     inputs = tokenizer([input_text], return_tensors="pt")
 
     outputs = model.generate(
         **inputs,
         max_new_tokens=number_steps,
         num_beams=number_beams,
-        num_return_sequences=number_beams,
+        num_return_sequences=number_sequences,
         return_dict_in_generate=True,
         length_penalty=length_penalty,
         output_scores=True,
@@ -443,6 +443,7 @@ Play with the parameters below to understand how beam search decoding works!
 - **Number of beams** (`num_beams`): the number of beams to use
 - **Length penalty** (`length_penalty`): the length penalty to apply to outputs. `length_penalty` > 0.0 promotes longer sequences, while `length_penalty` < 0.0 encourages shorter sequences.
 This parameter will not impact the beam search paths, but only influence the choice of sequences in the end towards longer or shorter sequences.
+- **Number of sequences** (`num_return_sequences`): the number of sequences to be returned at the end of generation.
 """
     )
     text = gr.Textbox(
@@ -450,21 +451,24 @@ This parameter will not impact the beam search paths, but only influence the cho
         value="Conclusion: thanks a lot. This article was originally published on",
     )
     with gr.Row():
-        steps = gr.Slider(
+        n_steps = gr.Slider(
             label="Number of steps", minimum=1, maximum=12, step=1, value=4
         )
-        beams = gr.Slider(
+        n_beams = gr.Slider(
             label="Number of beams", minimum=2, maximum=4, step=1, value=3
         )
         length_penalty = gr.Slider(
             label="Length penalty", minimum=-3, maximum=3, step=0.5, value=1
         )
+        n_sequences = gr.Slider(
+            label="Number of sequences", minimum=1, maximum=n_beams, step=1, value=n_beams
+        )
     button = gr.Button()
     out_html = gr.Markdown()
     out_markdown = gr.Markdown()
     button.click(
         get_beam_search_html,
-        inputs=[text, steps, beams, length_penalty],
+        inputs=[text, steps, beams, length_penalty, n_sequences],
         outputs=[out_html, out_markdown],
     )