Update app.py

app.py

@@ -272,7 +272,7 @@ class BeamNode:
     is_selected_sequence: bool
 
 
-def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
+def generate_beams(n_beams, start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
     original_tree = BeamNode(
         cumulative_score=0,
         current_token_ix=None,
@@ -284,7 +284,6 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
         is_final=False,
         is_selected_sequence=False,
     )
-    n_beams = len(scores[0])
     beam_trees = [original_tree] * n_beams
     generation_length = len(scores)
 
@@ -429,7 +428,7 @@ def get_beam_search_html(
     outputs = model.generate(
         **inputs,
         max_new_tokens=number_steps,
-        num_beams=number_beams,
+        num_beams=max(number_beams, 2),
         num_return_sequences=num_return_sequences,
         return_dict_in_generate=True,
         length_penalty=length_penalty,
@@ -447,6 +446,7 @@ def get_beam_search_html(
         markdown += f"\n- Score `{outputs.sequences_scores[i]:.2f}`: `{clean(sequence.replace('<s> ', ''))}`"
 
     original_tree = generate_beams(
+        number_beams,
         input_text,
         outputs.scores[:],
         length_penalty,
@@ -493,7 +493,7 @@ This parameter will not impact the beam search paths, but only influence the cho
             label="Number of steps", minimum=1, maximum=12, step=1, value=5
         )
         n_beams = gr.Slider(
-            label="Number of beams", minimum=2, maximum=4, step=1, value=4
+            label="Number of beams", minimum=1, maximum=4, step=1, value=4
         )
         length_penalty = gr.Slider(
             label="Length penalty", minimum=-3, maximum=3, step=0.5, value=1