Update app.py

app.py

@@ -300,26 +300,28 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
             token_scores,
         ) = ([], [], [], [], [], [])
 
-        for beam_ix in range(len(beam_trees)): # not range(n_beams) since there might be more ongoing trees.
+        score_idx = 0
+        for beam_ix in range(len(beam_trees)):
             current_beam = beam_trees[beam_ix]
 
             # skip if the beam is already final
             if current_beam.is_final:
                 continue
-
+                
             # Get top cumulative scores for the current beam
             current_top_token_indexes = list(
-                np.array(scores[step][beam_ix].argsort()[-n_beams:])[::-1]
+                np.array(scores[step][score_idx].argsort()[-n_beams:])[::-1]
             )
             top_token_indexes += current_top_token_indexes
-            token_scores += list(np.array(scores[step][beam_ix][current_top_token_indexes]))
+            token_scores += list(np.array(scores[step][score_idx][current_top_token_indexes]))
             top_cumulative_scores += list(
-                np.array(scores[step][beam_ix][current_top_token_indexes])
+                np.array(scores[step][score_idx][current_top_token_indexes])
                 + current_beam.cumulative_score
             )
             beam_indexes += [beam_ix] * n_beams
             current_sequence += [beam_trees[beam_ix].current_sequence] * n_beams
             top_tokens += [tokenizer.decode([el]) for el in current_top_token_indexes]
+            score_idx += 1
 
         top_df = pd.DataFrame.from_dict(
             {