Upload app.py

app.py

@@ -174,9 +174,6 @@ STYLE = """
 .nonselected-sequence {
     background-color: var(--primary-500);
 }
-.nopadding {
-    padding-left: 0!important;
-}
 """
 
 
@@ -241,7 +238,7 @@ def generate_nodes(node, step):
 def generate_html(start_sentence, original_tree):
     html_output = f"""<div class="custom-container">
 				<div class="tree">
-                <ul class="nopadding"> <li class="nopadding"> <a id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </a>"""
+                <ul> <li> <a href='#' id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </a>"""
     html_output += "<ul> "
     for subnode in original_tree.children.values():
         html_output += generate_nodes(subnode, step=1)
@@ -273,20 +270,20 @@ class BeamNode:
 
 
 def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
-    input_length = len(tokenizer([start_sentence], return_tensors="pt"))
     original_tree = BeamNode(
         cumulative_score=0,
         current_token_ix=None,
         table=None,
         current_sequence=start_sentence,
         children={},
-        children_score_divider=((input_length + 1) ** length_penalty),
+        children_score_divider=(1 ** length_penalty),
         total_score=None,
         is_final=False,
         is_selected_sequence=False,
     )
     n_beams = len(scores[0])
     beam_trees = [original_tree] * n_beams
+    generation_length = len(scores)
 
     for step, step_scores in enumerate(scores):
 
@@ -297,8 +294,11 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
             beam_indexes,
             current_sequence,
             top_tokens,
-        ) = ([], [], [], [], [])
-        for beam_ix in range(n_beams):
+            token_scores,
+        ) = ([], [], [], [], [], [])
+
+        score_idx = 0
+        for beam_ix in range(len(beam_trees)):
             current_beam = beam_trees[beam_ix]
 
             # skip if the beam is already final
@@ -307,16 +307,18 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
 
             # 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][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(
             {
@@ -325,6 +327,7 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
                 "beam_index": beam_indexes,
                 "current_sequence": current_sequence,
                 "token": top_tokens,
+                "token_score": token_scores,
             }
         )
         maxes = top_df.groupby(["token_index", "current_sequence"])[
@@ -333,78 +336,85 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
 
         top_df = top_df.loc[maxes]
 
-        # Sort all top probabilities and keep top n_beams
+        # Sort all top probabilities and keep top n_beams * 2 (* 2 because each beam may end this iteration, and we
+        # want to keep at least `n_beams` beams alive)
         top_df_selected = top_df.sort_values("cumulative_score", ascending=False).iloc[
-            :n_beams
+            :n_beams * 2
         ]
-        if any(["you enjoyed" in el for el in top_df["current_sequence"]]):
-            print("Displaying debug info:::")
-            print(top_df_selected)
+        beams_to_keep = 0
+        unfinished_beams = 0
+        for _, row in top_df_selected.iterrows():
+            beams_to_keep += 1
+            current_token_choice_ix = row["token_index"]
+            is_final = step == len(scores) - 1 or current_token_choice_ix == tokenizer.eos_token_id
+            if not is_final:
+                unfinished_beams += 1
+            if unfinished_beams >= n_beams:
+                break
+            if step == generation_length - 1 and beams_to_keep == n_beams:
+                break
+        top_df_selected_filtered = top_df_selected.iloc[:beams_to_keep]
 
         # Write the scores table - one per beam source
-        for beam_ix in reversed(list(range(n_beams))):
+        score_idx = 0
+        for beam_ix in range(len(beam_trees)):
             current_beam = beam_trees[beam_ix]
             if current_beam.table is None:
-                selected_tokens = top_df_selected.loc[
-                    top_df_selected["current_sequence"] == current_beam.current_sequence
+                selected_tokens = top_df_selected_filtered.loc[
+                    top_df_selected_filtered["current_sequence"] == current_beam.current_sequence
                 ]
                 markdown_table = generate_markdown_table(
-                    step_scores[beam_ix, :],
+                    step_scores[score_idx, :],
                     current_beam.cumulative_score,
                     current_beam.children_score_divider,
                     chosen_tokens=list(selected_tokens["token"].values),
                 )
                 beam_trees[beam_ix].table = markdown_table
+            if not current_beam.is_final:
+                score_idx = min(score_idx + 1, n_beams - 1)
 
         # Add new children to each beam
         cumulative_scores = [beam.cumulative_score for beam in beam_trees]
-        for _, row in top_df_selected.iterrows():
+        for _, row in top_df_selected_filtered.iterrows():
             # Update the source tree
             source_beam_ix = int(row["beam_index"])
             current_token_choice_ix = row["token_index"]
             current_token_choice = tokenizer.decode([current_token_choice_ix])
+            token_scores = row["token_score"]
 
-            cumulative_score = (
-                cumulative_scores[source_beam_ix]
-                + scores[step][source_beam_ix][current_token_choice_ix].numpy()
-            )
+            cumulative_score = cumulative_scores[source_beam_ix] + np.asarray(token_scores)
             current_sequence = (
                 beam_trees[source_beam_ix].current_sequence + current_token_choice
             )
-            if current_token_choice_ix == 340:
-                print("Found info:")
-                print(f"We generate token '{current_token_choice}', and the total sequence is '{current_sequence}'")
+            is_final = step == len(scores) - 1 or current_token_choice_ix == tokenizer.eos_token_id
             beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
                 current_token_ix=current_token_choice_ix,
                 table=None,
                 children={},
                 current_sequence=current_sequence,
                 cumulative_score=cumulative_score,
-                total_score=cumulative_score
-                / ((input_length + step - 1) ** length_penalty),
-                children_score_divider=((input_length + step) ** length_penalty),
-                is_final=(
-                    step == len(scores) - 1
-                    or current_token_choice_ix == tokenizer.eos_token_id
-                ),
+                total_score=cumulative_score / (step + 1 ** length_penalty),
+                children_score_divider=((step + 2) ** length_penalty),
+                is_final=is_final,
                 is_selected_sequence=(
                     current_sequence.replace("<|endoftext|>", "")
                     in [el.replace("<|endoftext|>", "") for el in decoded_sequences]
                 ),
             )
 
-
         # Swap all beams by descending cumul score, so that n°1 has the highest cumulative score, and so on
         beam_trees = [
-            beam_trees[int(top_df_selected.iloc[beam_ix]["beam_index"])]
-            for beam_ix in range(n_beams)
+            beam_trees[int(top_df_selected_filtered.iloc[beam_ix]["beam_index"])]
+            for beam_ix in range(beams_to_keep)
         ]
 
         # Advance all beams by one token
-        for beam_ix in range(n_beams):
-            current_token_choice_ix = top_df_selected.iloc[beam_ix]["token_index"]
+        for beam_ix in range(beams_to_keep):
+            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
@@ -459,9 +469,9 @@ with gr.Blocks(
 ) as demo:
     gr.Markdown(
         """# <span style='color:var(--primary-500)!important'>Beam Search Visualizer</span>
-                
+
 Play with the parameters below to understand how beam search decoding works!
-                
+
 #### <span style='color:var(--primary-500)!important'>Parameters:</span>
 - **Sentence to decode from** (`inputs`): the input sequence to your decoder.
 - **Number of steps** (`max_new_tokens`): the number of tokens to generate.
@@ -473,20 +483,20 @@ This parameter will not impact the beam search paths, but only influence the cho
     )
     text = gr.Textbox(
         label="Sentence to decode from",
-        value="Thank you for",
+        value="Conclusion: thanks a lot. That's all for today",
     )
     with gr.Row():
         n_steps = gr.Slider(
-            label="Number of steps", minimum=1, maximum=10, step=1, value=4
+            label="Number of steps", minimum=1, maximum=10, step=1, value=10
         )
         n_beams = gr.Slider(
-            label="Number of beams", minimum=2, maximum=4, step=1, value=3
+            label="Number of beams", minimum=2, maximum=4, step=1, value=4
         )
         length_penalty = gr.Slider(
             label="Length penalty", minimum=-3, maximum=3, step=0.5, value=1
         )
         num_return_sequences = gr.Slider(
-            label="Number of return sequences", minimum=1, maximum=3, step=1, value=2
+            label="Number of return sequences", minimum=1, maximum=4, step=1, value=3
         )
 
     n_beams.change(
@@ -501,4 +511,4 @@ This parameter will not impact the beam search paths, but only influence the cho
         outputs=[out_html, out_markdown],
     )
 
-demo.launch()
+demo.launch()