Update app.py

app.py

@@ -130,7 +130,7 @@ STYLE = """
     margin-top: -5px;
     transform: rotate(315deg);
 }
-.box {
+.tree li a {
 	border: 1px solid var(--body-text-color);
 	padding: 5px;
 	border-radius: 5px;
@@ -141,19 +141,18 @@ STYLE = """
     align-items: center;
     justify-content: space-between;
     overflow: hidden;
-    cursor: pointer;
 }
-.box span {
+.tree li a span {
 	padding: 5px;
 	font-size: 12px;
 	letter-spacing: 1px;
 	font-weight: 500;
 }
 /*Hover-Section*/
-.box:hover, .box:hover+ul li .box {
+.tree li a:hover, .tree li a:hover+ul li a {
 	background: var(--primary-500);
 }
-.box:hover+ul li::after, .box:hover+ul li::before, .box:hover+ul::before, .box:hover+ul ul::before, .box:hover+ul .box::before {
+.tree li a:hover+ul li::after, .tree li a:hover+ul li::before, .tree li a:hover+ul::before, .tree li a:hover+ul ul::before, .tree li a:hover+ul a::before {
 	border-color: var(--primary-500);
 }
 .chosen-token {
@@ -175,9 +174,6 @@ STYLE = """
 .nonselected-sequence {
     background-color: var(--primary-500);
 }
-.nomargin {
-    padding-left: 0!important;
-}
 """
 
 
@@ -220,14 +216,14 @@ def generate_nodes(node, step):
             selected_class = "selected-sequence"
         else:
             selected_class = "nonselected-sequence"
-        return f"<li> <div class='box end-of-text child {selected_class}'> <span> <b>{clean(token)}</b> <br>Total score: {node.total_score:.2f}</span> </div> </li>"
+        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> <div class='box nonfinal child'> <span> <b>{clean(token)}</b> </span>"
+        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 += "</div>"
+    html_content += "</a>"
 
     if len(node.children.keys()) > 0:
         html_content += "<ul> "
@@ -241,15 +237,16 @@ def generate_nodes(node, step):
 
 def generate_html(start_sentence, original_tree):
     html_output = f"""<div class="custom-container">
-				<div class="tree"> <ul class="nomargin"><li class="nomargin">
-                <div class="box" id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </div>"""
+				<div class="tree">
+                <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)
     html_output += "</ul>"
     html_output += """
-        </li></ul></div>
-    </div>
+        </li> </ul>
+        </div>
+    </body>
     """
     return html_output
 
@@ -272,7 +269,7 @@ class BeamNode:
     is_selected_sequence: bool
 
 
-def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
+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,
@@ -289,6 +286,8 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
     beam_trees = [original_tree] * n_beams
 
     for step, step_scores in enumerate(scores):
+
+        # Gather all possible descendants for each beam
         (
             top_token_indexes,
             top_cumulative_scores,
@@ -296,7 +295,7 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
             current_sequence,
             top_tokens,
         ) = ([], [], [], [], [])
-        for beam_ix in range(n_beams):  # Get possible descendants for each beam
+        for beam_ix in range(n_beams):
             current_beam = beam_trees[beam_ix]
 
             # skip if the beam is already final
@@ -316,7 +315,6 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
             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,
@@ -336,6 +334,9 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
         top_df_selected = top_df.sort_values("cumulative_score", ascending=False).iloc[
             :n_beams
         ]
+        if any(["you enjoyed" in el for el in top_df["current_sequence"]]):
+            print("Displaying debug info:::")
+            display(top_df_selected)
 
         # Write the scores table - one per beam source
         for beam_ix in reversed(list(range(n_beams))):
@@ -352,14 +353,13 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
                 )
                 beam_trees[beam_ix].table = markdown_table
 
-        # Add new children for each beam
+        # Add new children to each beam
         cumulative_scores = [beam.cumulative_score for beam in beam_trees]
-        for beam_ix in range(n_beams):
-            current_token_choice_ix = top_df_selected.iloc[beam_ix]["token_index"]
-            current_token_choice = tokenizer.decode([current_token_choice_ix])
-
+        for _, row in top_df_selected.iterrows():
             # Update the source tree
-            source_beam_ix = int(top_df_selected.iloc[beam_ix]["beam_index"])
+            source_beam_ix = int(row["beam_index"])
+            current_token_choice_ix = row["token_index"]
+            current_token_choice = tokenizer.decode([current_token_choice_ix])
 
             cumulative_score = (
                 cumulative_scores[source_beam_ix]
@@ -368,6 +368,9 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
             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}'")
             beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
                 current_token_ix=current_token_choice_ix,
                 table=None,
@@ -387,7 +390,8 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
                 ),
             )
 
-        # Reassign all beams at once
+
+        # 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)
@@ -400,7 +404,6 @@ 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
@@ -432,6 +435,7 @@ def get_beam_search_html(
         outputs.scores[:],
         length_penalty,
         decoded_sequences,
+        outputs.beam_indices,
     )
     html = generate_html(input_text, original_tree)
     return html, markdown
@@ -466,7 +470,7 @@ This parameter will not impact the beam search paths, but only influence the cho
     )
     text = gr.Textbox(
         label="Sentence to decode from",
-        value="Conclusion: thanks a lot. This article was originally published on",
+        value="Conclusion: thanks a lot. That's all for today",
     )
     with gr.Row():
         n_steps = gr.Slider(