Update app.py

app.py

@@ -7,7 +7,6 @@ import spaces
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 model = AutoModelForCausalLM.from_pretrained("gpt2")
 
-tokenizer.pad_token_id = tokenizer.eos_token_id
 print("Loading finished.")
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -174,22 +173,13 @@ a:before {
 """
 
 
-def generate_nodes(token, node, step=0):
-    """Recursively generate HTML for the tree nodes."""
-
-    html_content = f" <li> <a href='#' class='{('chosen' if node.table is None else '')}' id='{('root' if step==0 else '')}'> <span> <b>{token}</b> </span> "
-    html_content += node.table if node.table is not None else ""
-    html_content += "</a>"
-    if len(node.children.keys()) > 0:
-        html_content += "<ul> "
-        for token, subnode in node.children.items():
-            html_content += generate_nodes(token, subnode, step=step + 1)
-        html_content += "</ul>"
-    html_content += "</li>"
-    return html_content
+def clean(s):
+    return s.replace("\n", r"\n").replace("\t", r"\t")
 
 
-def generate_markdown_table(scores, sequence_prob, top_k=4, chosen_tokens=None):
+def generate_markdown_table(
+    scores, previous_cumul_score, score_divider, top_k=4, chosen_tokens=None
+):
     markdown_table = """
     <table>
         <tr>
@@ -204,21 +194,41 @@ def generate_markdown_table(scores, sequence_prob, top_k=4, chosen_tokens=None):
             item_class = "chosen"
         markdown_table += f"""
         <tr class={item_class}>
-            <td>{token}</td>
+            <td>{clean(token)}</td>
             <td>{scores[token_idx]:.4f}</td>
-            <td>{scores[token_idx] + sequence_prob:.4f}</td>
+            <td>{(scores[token_idx] + previous_cumul_score)/score_divider:.4f}</td>
         </tr>"""
     markdown_table += """
     </table>"""
     return markdown_table
 
 
+def generate_nodes(token_ix, node, step):
+    """Recursively generate HTML for the tree nodes."""
+    token = tokenizer.decode([token_ix])
+    html_content = f" <li> <a href='#' class='{('chosen' if node.table is None else '')}'> <span> <b>{token_ix}:<br>{clean(token)}</b> </span> "
+    html_content += node.table if node.table is not None else ""
+    html_content += "</a>"
+    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 += "</ul>"
+    html_content += "</li>"
+    return html_content
+
+
 def generate_html(start_sentence, original_tree):
 
-    html_output = """<div class="custom-container">
+    html_output = f"""<div class="custom-container">
 				<div class="tree">
-                <ul>"""
-    html_output += generate_nodes(start_sentence, original_tree)
+                <ul>
+                <li> <a href='#' id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </a>"""
+    if len(original_tree.children.keys()) > 0:
+        html_output += "<ul> "
+        for token_ix, subnode in original_tree.children.items():
+            html_output += generate_nodes(token_ix, subnode, step=1)
+        html_output += "</ul>"
 
     html_output += """
         </ul>
@@ -236,16 +246,21 @@ from dataclasses import dataclass
 @dataclass
 class BeamNode:
     cumulative_score: float
+    children_score_divider: float
     table: str
     current_sentence: str
-    children: Dict[str, "BeamNode"]
+    children: Dict[int, "BeamNode"]
 
 
-def generate_beams(start_sentence, scores, sequences, beam_indices):
-    print(tokenizer.batch_decode(sequences))
+def generate_beams(start_sentence, scores, sequences, length_penalty):
     sequences = sequences.cpu().numpy()
+    input_length = len(tokenizer([start_sentence], return_tensors="pt"))
     original_tree = BeamNode(
-        cumulative_score=0, table=None, current_sentence=start_sentence, children={}
+        cumulative_score=0,
+        table=None,
+        current_sentence=start_sentence,
+        children={},
+        children_score_divider=((input_length + 1) ** length_penalty),
     )
     n_beams = len(scores[0])
     beam_trees = [original_tree] * n_beams
@@ -302,6 +317,7 @@ def generate_beams(start_sentence, scores, sequences, beam_indices):
             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
@@ -315,18 +331,18 @@ def generate_beams(start_sentence, scores, sequences, beam_indices):
             # Update the source tree
             source_beam_ix = int(top_df_selected.iloc[beam_ix]["beam_index"])
 
-            previous_len = len(str(original_tree))
-            beam_trees[source_beam_ix].children[current_token_choice] = BeamNode(
+            cumulative_score = (
+                cumulative_scores[source_beam_ix]
+                + scores[step][source_beam_ix][current_token_choice_ix].numpy()
+            )
+            beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
                 table=None,
                 children={},
                 current_sentence=beam_trees[source_beam_ix].current_sentence
                 + current_token_choice,
-                cumulative_score=cumulative_scores[source_beam_ix]
-                + scores[step][source_beam_ix][current_token_choice_ix].numpy(),
+                cumulative_score=cumulative_score,
+                children_score_divider=((input_length + step + 1) ** length_penalty),
             )
-            assert (
-                len(str(original_tree)) > previous_len
-            ), "Original tree has not increased size"
 
         # Reassign all beams at once
         beam_trees = [
@@ -337,12 +353,12 @@ def generate_beams(start_sentence, scores, sequences, beam_indices):
         # 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"]
-            current_token_choice = tokenizer.decode([current_token_choice_ix])
-            beam_trees[beam_ix] = beam_trees[beam_ix].children[current_token_choice]
+            beam_trees[beam_ix] = beam_trees[beam_ix].children[current_token_choice_ix]
+
     return original_tree
 
 @spaces.GPU
-def get_beam_search_html(input_text, number_steps, number_beams):
+def get_beam_search_html(input_text, number_steps, number_beams, length_penalty):
     inputs = tokenizer([input_text], return_tensors="pt")
 
     outputs = model.generate(
@@ -351,19 +367,21 @@ def get_beam_search_html(input_text, number_steps, number_beams):
         num_beams=number_beams,
         num_return_sequences=number_beams,
         return_dict_in_generate=True,
+        length_penalty=-10.0,
         output_scores=True,
-        top_k=5,
         do_sample=False,
     )
+    print("Sequences:")
+    print(tokenizer.batch_decode(outputs.sequences))
+    print("Scores:", outputs.sequences_scores)
 
     original_tree = generate_beams(
         input_text,
         outputs.scores[:],
         outputs.sequences[:, :],
-        outputs.beam_indices[:, :],
+        length_penalty,
     )
     html = generate_html(input_text, original_tree)
-    print(html)
     return html
 
 
@@ -374,10 +392,12 @@ with gr.Blocks(
     css=STYLE,
 ) as demo:
     text = gr.Textbox(label="Sentence to decode from", value="Today is")
-    steps = gr.Slider(label="Number of steps", minimum=1, maximum=8, step=1, value=4)
-    beams = gr.Slider(label="Number of beams", minimum=2, maximum=4, step=1, value=3)
+    with gr.Row():
+        steps = gr.Slider(label="Number of steps", minimum=1, maximum=8, step=1, value=4)
+        beams = gr.Slider(label="Number of beams", minimum=2, maximum=4, step=1, value=3)
+        length_penalty = gr.Slider(label="Length penalty", minimum=-5, maximum=5, step=0.5, value=1)
     button = gr.Button()
     out = gr.Markdown(label="Output")
-    button.click(get_beam_search_html, inputs=[text, steps, beams], outputs=out)
+    button.click(get_beam_search_html, inputs=[text, steps, beams, length_penalty], outputs=out)
 
 demo.launch()