Update app.py

app.py

@@ -218,7 +218,9 @@ def generate_nodes(node, step):
             selected_class = "nonselected-sequence"
         return f"<li> <a href='#' class='end-of-text child {selected_class}'> <span> <b>{clean(token)}</b> <br>Total score: {node.total_score:.2f}</span> </a> </li>"
 
-    html_content = f"<li> <a href='#' class='nonfinal child'> <span> <b>{clean(token)}</b> </span>"
+    html_content = (
+        f"<li> <a href='#' class='nonfinal child'> <span> <b>{clean(token)}</b> </span>"
+    )
     if node.table is not None:
         html_content += node.table
     html_content += "</a>"
@@ -288,7 +290,7 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
             top_token_indexes,
             top_cumulative_scores,
             beam_indexes,
-            current_completions,
+            current_sequence,
             top_tokens,
         ) = ([], [], [], [], [])
         for beam_ix in range(n_beams):  # Get possible descendants for each beam
@@ -308,19 +310,20 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
                 + current_beam.cumulative_score
             )
             beam_indexes += [beam_ix] * n_beams
-            current_completions += [beam_trees[beam_ix].current_sequence] * n_beams
+            current_sequence += [beam_trees[beam_ix].current_sequence] * n_beams
             top_tokens += [tokenizer.decode([el]) for el in current_top_token_indexes]
 
+
         top_df = pd.DataFrame.from_dict(
             {
                 "token_index": top_token_indexes,
                 "cumulative_score": top_cumulative_scores,
                 "beam_index": beam_indexes,
-                "current_completions": current_completions,
+                "current_sequence": current_sequence,
                 "token": top_tokens,
             }
         )
-        maxes = top_df.groupby(["token_index", "current_completions"])[
+        maxes = top_df.groupby(["token_index", "current_sequence"])[
             "cumulative_score"
         ].idxmax()
 
@@ -331,20 +334,20 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
             :n_beams
         ]
 
-        # Write the scores table - one per beam source?
-        # Edge case: if several beam indexes are actually on the same beam, the selected tokens by beam_index for the second one will be empty. So we reverse
+        # Write the scores table - one per beam source
         for beam_ix in reversed(list(range(n_beams))):
             current_beam = beam_trees[beam_ix]
-            selected_tokens = top_df_selected.loc[
-                top_df_selected["beam_index"] == beam_ix
-            ]
-            markdown_table = generate_markdown_table(
-                step_scores[beam_ix, :],
-                current_beam.cumulative_score,
-                current_beam.children_score_divider,
-                chosen_tokens=list(selected_tokens["token"].values),
-            )
-            beam_trees[beam_ix].table = markdown_table
+            if current_beam.table is None:
+                selected_tokens = top_df_selected.loc[
+                    top_df_selected["current_sequence"] == current_beam.current_sequence
+                ]
+                markdown_table = generate_markdown_table(
+                    step_scores[beam_ix, :],
+                    current_beam.cumulative_score,
+                    current_beam.children_score_divider,
+                    chosen_tokens=list(selected_tokens["token"].values),
+                )
+                beam_trees[beam_ix].table = markdown_table
 
         # Add new children for each beam
         cumulative_scores = [beam.cumulative_score for beam in beam_trees]
@@ -375,7 +378,10 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
                     step == len(scores) - 1
                     or current_token_choice_ix == tokenizer.eos_token_id
                 ),
-                is_selected_sequence=(current_sequence.replace('<|endoftext|>', '') in [el.replace('<|endoftext|>', '') for el in decoded_sequences]),
+                is_selected_sequence=(
+                    current_sequence.replace("<|endoftext|>", "")
+                    in [el.replace("<|endoftext|>", "") for el in decoded_sequences]
+                ),
             )
 
         # Reassign all beams at once
@@ -391,8 +397,11 @@ 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, num_return_sequences):
+def get_beam_search_html(
+    input_text, number_steps, number_beams, length_penalty, num_return_sequences
+):
     inputs = tokenizer([input_text], return_tensors="pt")
 
     outputs = model.generate(
@@ -426,7 +435,9 @@ def get_beam_search_html(input_text, number_steps, number_beams, length_penalty,
 
 
 def change_num_return_sequences(n_beams):
-    return gr.Slider(label="Number of sequences", minimum=1, maximum=n_beams, step=1, value=n_beams)
+    return gr.Slider(
+        label="Number of sequences", minimum=1, maximum=n_beams, step=1, value=n_beams
+    )
 
 
 with gr.Blocks(
@@ -468,7 +479,9 @@ This parameter will not impact the beam search paths, but only influence the cho
             label="Number of return sequences", minimum=1, maximum=3, step=1, value=2
         )
 
-    n_beams.change(fn=change_num_return_sequences, inputs=n_beams, outputs=num_return_sequences)
+    n_beams.change(
+        fn=change_num_return_sequences, inputs=n_beams, outputs=num_return_sequences
+    )
     button = gr.Button()
     out_html = gr.Markdown()
     out_markdown = gr.Markdown()