Update app.py

app.py

@@ -2,7 +2,6 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import numpy as np
 import gradio as gr
-import spaces
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 model = AutoModelForCausalLM.from_pretrained("gpt2")
@@ -118,7 +117,7 @@ STYLE = """
 .tree ul:has(> li:only-child)::before {
     width:40px;
 }
-.tree li a:before {
+.child:before {
     border-right: 2px solid var(--body-text-color);
     border-bottom: 2px solid var(--body-text-color);
     content: "";
@@ -150,13 +149,13 @@ STYLE = """
 }
 /*Hover-Section*/
 .tree li a:hover, .tree li a:hover+ul li a {
-	background: #ffedd5;
+	background: var(--primary-700);
 }
-.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 {
-	border-color: #7c2d12;
+.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-200);
 }
-.end-of-text, .chosen {
-    background-color: #ea580c;
+.end-of-text, .chosen-token {
+    background-color: var(--primary-600);
 }
 .end-of-text {
     width:auto!important;
@@ -164,7 +163,10 @@ STYLE = """
 .nonfinal {
     width:280px;
     min-width: 280px;
-}   
+}
+.selected-sequence {
+    background-color: var(--secondary-600)!important;
+}
 """
 
 
@@ -186,7 +188,7 @@ def generate_markdown_table(
         token = tokenizer.decode([token_idx])
         item_class = ""
         if chosen_tokens and token in chosen_tokens:
-            item_class = "chosen"
+            item_class = "chosen-token"
         markdown_table += f"""
         <tr class={item_class}>
             <td>{clean(token)}</td>
@@ -198,16 +200,16 @@ def generate_markdown_table(
     return markdown_table
 
 
-def generate_nodes(token_ix, node, step):
+def generate_nodes(node, step):
     """Recursively generate HTML for the tree nodes."""
-    token = tokenizer.decode([token_ix])
+    token = tokenizer.decode([node.current_token_ix])
+
+    selected_class = "selected-sequence" if node.is_selected_sequence else ""
 
     if node.is_final:
-        return f"<li> <a href='#' class='end-of-text'> <span> <b>{clean(token)}</b> <br>Total score: {node.total_score:.2f}</span> </a> </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> <a href='#' class='nonfinal'> <span> <b>{clean(token)}</b> </span>"
-    )
+    html_content = f"<li> <a href='#' class='nonfinal child {selected_class}'> <span> <b>{clean(token)}</b> </span>"
     if node.table is not None:
         html_content += node.table
     html_content += "</a>"
@@ -215,7 +217,7 @@ def generate_nodes(token_ix, node, step):
     if len(node.children.keys()) > 0:
         html_content += "<ul> "
         for token_ix, subnode in node.children.items():
-            html_content += generate_nodes(token_ix, subnode, step=step + 1)
+            html_content += generate_nodes(subnode, step=step + 1)
         html_content += "</ul>"
     html_content += "</li>"
 
@@ -227,8 +229,8 @@ def generate_html(start_sentence, original_tree):
 				<div class="tree">
                 <ul> <li> <a href='#' id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </a>"""
     html_output += "<ul> "
-    for token_ix, subnode in original_tree.children.items():
-        html_output += generate_nodes(token_ix, subnode, step=1)
+    for subnode in original_tree.children.values():
+        html_output += generate_nodes(subnode, step=1)
     html_output += "</ul>"
     html_output += """
         </li> </ul>
@@ -249,24 +251,25 @@ class BeamNode:
     cumulative_score: float
     children_score_divider: float
     table: str
-    current_sentence: str
+    current_sequence: str
     children: Dict[int, "BeamNode"]
     total_score: float
     is_final: bool
+    is_selected_sequence: bool
 
 
-def generate_beams(start_sentence, scores, sequences, length_penalty):
-    sequences = sequences.cpu().numpy()
+def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
     input_length = len(tokenizer([start_sentence], return_tensors="pt"))
     original_tree = BeamNode(
         cumulative_score=0,
         current_token_ix=None,
         table=None,
-        current_sentence=start_sentence,
+        current_sequence=start_sentence,
         children={},
         children_score_divider=((input_length + 1) ** length_penalty),
         total_score=None,
         is_final=False,
+        is_selected_sequence=False,
     )
     n_beams = len(scores[0])
     beam_trees = [original_tree] * n_beams
@@ -296,7 +299,7 @@ def generate_beams(start_sentence, scores, sequences, length_penalty):
                 + current_beam.cumulative_score
             )
             beam_indexes += [beam_ix] * n_beams
-            current_completions += [beam_trees[beam_ix].current_sentence] * n_beams
+            current_completions += [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(
@@ -347,12 +350,14 @@ def generate_beams(start_sentence, scores, sequences, length_penalty):
                 cumulative_scores[source_beam_ix]
                 + scores[step][source_beam_ix][current_token_choice_ix].numpy()
             )
+            current_sequence = (
+                beam_trees[source_beam_ix].current_sequence + current_token_choice
+            )
             beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
                 current_token_ix=current_token_choice_ix,
                 table=None,
                 children={},
-                current_sentence=beam_trees[source_beam_ix].current_sentence
-                + current_token_choice,
+                current_sequence=current_sequence,
                 cumulative_score=cumulative_score,
                 total_score=cumulative_score
                 / ((input_length + step - 1) ** length_penalty),
@@ -361,6 +366,7 @@ def generate_beams(start_sentence, scores, sequences, length_penalty):
                     step == len(scores) - 1
                     or current_token_choice_ix == tokenizer.eos_token_id
                 ),
+                is_selected_sequence=(current_sequence in decoded_sequences),
             )
 
         # Reassign all beams at once
@@ -376,7 +382,7 @@ def generate_beams(start_sentence, scores, sequences, length_penalty):
 
     return original_tree
 
-@spaces.GPU
+
 def get_beam_search_html(input_text, number_steps, number_beams, length_penalty):
     inputs = tokenizer([input_text], return_tensors="pt")
 
@@ -390,17 +396,21 @@ def get_beam_search_html(input_text, number_steps, number_beams, length_penalty)
         output_scores=True,
         do_sample=False,
     )
-    markdown = "Output sequences:"
+    markdown = "The conclusive sequences are the ones that end in an `<|endoftext|>` token or at the end of generation."
+    markdown += "\n\nThey are ranked by their scores, as given by the formula `score = cumulative_score / (output_length ** length_penalty)`.\n\n"
+    markdown += "Only the top `num_beams` scoring sequences are returned: in the tree they are highlighted in **<span style='color:var(--secondary-600)!important'>blue</span>**."
+    markdown += " The non-selected sequences are also shown in the tree, highlighted in **<span style='color:var(--primary-600)!important'>yellow</span>**."
+    markdown += "\n#### <span style='color:var(--secondary-600)!important'>Output sequences:</span>"
     # Sequences are padded anyway so you can batch decode them
     decoded_sequences = tokenizer.batch_decode(outputs.sequences)
     for i, sequence in enumerate(decoded_sequences):
-        markdown += f"\n- '{clean(sequence.replace('<s> ', ''))}' (score {outputs.sequences_scores[i]:.2f})"
+        markdown += f"\n- Score `{outputs.sequences_scores[i]:.2f}`: `{clean(sequence.replace('<s> ', ''))}`"
 
     original_tree = generate_beams(
         input_text,
         outputs.scores[:],
-        outputs.sequences[:, :],
         length_penalty,
+        decoded_sequences,
     )
     html = generate_html(input_text, original_tree)
     return html, markdown
@@ -408,20 +418,21 @@ def get_beam_search_html(input_text, number_steps, number_beams, length_penalty)
 
 with gr.Blocks(
     theme=gr.themes.Soft(
-        text_size="lg", font=["monospace"], primary_hue=gr.themes.colors.yellow
+        primary_hue=gr.themes.colors.yellow,
+        secondary_hue=gr.themes.colors.blue,
     ),
     css=STYLE,
 ) as demo:
     gr.Markdown(
-        """# Beam search visualizer
+        """# <span style='color:var(--primary-600)!important'>Beam Search Visualizer</span>
                 
 Play with the parameters below to understand how beam search decoding works!
                 
-#### Parameters: 
-- **Sentence to decode from**: the input sequence to your decoder.
-- **Number of steps**: the number of tokens to generate
-- **Number of beams**: the number of beams to use
-- **Length penalty**: the length penalty to apply to outputs. `length_penalty` > 0.0 promotes longer sequences, while `length_penalty` < 0.0 encourages shorter sequences.
+#### <span style='color:var(--primary-600)!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
+- **Number of beams** (`num_beams`): the number of beams to use
+- **Length penalty** (`length_penalty`): the length penalty to apply to outputs. `length_penalty` > 0.0 promotes longer sequences, while `length_penalty` < 0.0 encourages shorter sequences.
 This parameter will not impact the beam search paths, but only influence the choice of sequences in the end towards longer or shorter sequences.
 """
     )