Update app.py

app.py

@@ -272,7 +272,7 @@ class BeamNode:
     is_selected_sequence: bool
 
 
-def generate_beams(n_beams, start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
+def generate_beams(n_beams, start_sentence, scores, length_penalty, decoded_sequences):
     original_tree = BeamNode(
         cumulative_score=0,
         current_token_ix=None,
@@ -415,8 +415,6 @@ def generate_beams(n_beams, start_sentence, scores, length_penalty, decoded_sequ
             current_token_choice_ix = top_df_selected_filtered.iloc[beam_ix]["token_index"]
             beam_trees[beam_ix] = beam_trees[beam_ix].children[current_token_choice_ix]
 
-        print(f"Step {step}, beams kept: {beams_to_keep}")
-
     return original_tree
 
 @spaces.GPU
@@ -445,14 +443,23 @@ def get_beam_search_html(
     for i, sequence in enumerate(decoded_sequences):
         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,
-        decoded_sequences,
-        outputs.beam_indices,
-    )
+    if number_beams > 1:
+        original_tree = generate_beams(
+            number_beams,
+            input_text,
+            outputs.scores[:],
+            length_penalty,
+            decoded_sequences,
+        )
+    else:
+        original_tree = generate_beams(
+            n_beams,
+            start_sentence,
+            outputs.logits,
+            0,
+            decoded_sequences,
+        )
+        
     html = generate_html(input_text, original_tree)
     return html, markdown