Update app.py

app.py

@@ -4,28 +4,43 @@ import numpy as np
 import gradio as gr
 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()}")
 # True
 if torch.cuda.is_available():
     print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
 STYLE = """
-.container {
+.custom-container {
 	width: 100%;
 	display: grid;
 	align-items: center;
     margin: 0!important;
+    overflow: scroll;
 }
 .prose ul ul {
     margin: 0!important;
-    font-size: 13px!important;
+    font-size: 10px!important;
 }
+.prose td, th {
+    padding-left: 2px;
+    padding-right: 2px;
+    padding-top: 0;
+    padding-bottom: 0;
+}
+
 .tree {
 	padding: 0px;
 	margin: 0!important;
 	box-sizing: border-box;
-    font-size: 16px;
+    font-size: 10px;
 	width: 100%;
+    min-width: 2000px;
 	height: auto;
 	text-align: center;
 }
@@ -34,13 +49,17 @@ STYLE = """
 	position: relative;
 	transition: .5s;
     margin: 0!important;
+    display: flex;
+    flex-direction: row;
+    justify-content: center;
+    gap:10px;
 }
 .tree li {
 	display: inline-table;
 	text-align: center;
 	list-style-type: none;
 	position: relative;
-	padding: 10px;
+	padding-top: 10px;
 	transition: .5s;
 }
 .tree li::before, .tree li::after {
@@ -88,7 +107,7 @@ STYLE = """
 }
 .tree li a {
 	border: 1px solid #ccc;
-	padding: 10px;
+	padding: 5px;
 	display: inline-grid;
 	border-radius: 5px;
 	text-decoration-line: none;
@@ -96,10 +115,8 @@ STYLE = """
 	transition: .5s;
 }
 .tree li a span {
-	border: 1px solid #ccc;
-	border-radius: 5px;
 	color: #666;
-	padding: 8px;
+	padding: 5px;
 	font-size: 12px;
 	text-transform: uppercase;
 	letter-spacing: 1px;
@@ -109,30 +126,26 @@ STYLE = """
 .tree li a:hover, .tree li a:hover i, .tree li a:hover span, .tree li a:hover+ul li a {
 	background: #c8e4f8;
 	color: #000;
-	border: 1px solid #94a0b4;
 }
 .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: #94a0b4;
 }
+.chosen {
+    background-color: red;
+}
 """
 
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-model = AutoModelForCausalLM.from_pretrained("gpt2")
-
-tokenizer.pad_token_id = tokenizer.eos_token_id
-print("Loading finished.")
-
 
-def generate_html(token, node):
-    """Recursively generate HTML for the tree."""
+def generate_nodes(token, node):
+    """Recursively generate HTML for the tree nodes."""
 
-    html_content = f" <li> <a href='#'> <span> <b>{token}</b> </span> "
-    html_content += node["table"] if node["table"] is not None else ""
+    html_content = f" <li> <a href='#' class={('chosen' if node.table is None 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:
+    if len(node.children.keys()) > 0:
         html_content += "<ul> "
-        for token, subnode in node["children"].items():
-            html_content += generate_html(token, subnode)
+        for token, subnode in node.children.items():
+            html_content += generate_nodes(token, subnode)
         html_content += "</ul>"
     html_content += "</li>"
     return html_content
@@ -144,9 +157,9 @@ def generate_markdown_table(scores, sequence_prob, top_k=4, chosen_tokens=None):
         <tr>
             <th><b>Token</b></th>
             <th><b>Step score</b></th>
-            <th><b>Cumulative score</b></th>
+            <th><b>Total score</b></th>
         </tr>"""
-    for token_idx in np.argsort(scores)[-top_k:]:
+    for token_idx in np.array(np.argsort(scores)[-top_k:])[::-1]:
         token = tokenizer.decode([token_idx])
         style = ""
         if chosen_tokens and token in chosen_tokens:
@@ -162,50 +175,140 @@ def generate_markdown_table(scores, sequence_prob, top_k=4, chosen_tokens=None):
     return markdown_table
 
 
-def display_tree(start_sentence, scores, sequences, beam_indices):
-    display = """<div class="container">
+def generate_html(start_sentence, original_tree):
+
+    html_output = """<div class="custom-container">
 				<div class="tree">
                 <ul>"""
-    sequences = sequences.cpu().numpy()
+    html_output += generate_nodes(start_sentence, original_tree)
+
+    html_output += """
+        </ul>
+        </div>
+    </body>
+    """
+    return html_output
+
+
+import pandas as pd
+from typing import Dict
+from dataclasses import dataclass
+
+
+@dataclass
+class BeamNode:
+    cumulative_score: float
+    table: str
+    current_sentence: str
+    children: Dict[str, "BeamNode"]
+
+
+def generate_beams(start_sentence, scores, sequences, beam_indices):
     print(tokenizer.batch_decode(sequences))
-    original_tree = {"table": None, "children": {}}
-    for sequence_ix in range(len(sequences)):
-        current_sequence_score = 0
-        current_tree = original_tree
-        for step, step_scores in enumerate(scores):
-            current_token_choice_ix = sequences[sequence_ix, step]
-            current_token_choice = tokenizer.decode([current_token_choice_ix])
-            current_beam = beam_indices[sequence_ix, step]
+    sequences = sequences.cpu().numpy()
+    original_tree = BeamNode(
+        cumulative_score=0, table=None, current_sentence=start_sentence, children={}
+    )
+    n_beams = len(scores[0])
+    beam_trees = [original_tree] * n_beams
+    for step, step_scores in enumerate(scores):
+        (
+            top_token_indexes,
+            top_cumulative_scores,
+            beam_indexes,
+            current_completions,
+            top_tokens,
+        ) = ([], [], [], [], [])
+        for beam_ix in range(n_beams):
+            current_beam = beam_trees[beam_ix]
+            # Get top cumulative scores for the current beam
+            current_top_token_indexes = list(
+                np.array(scores[step][beam_ix].argsort()[-n_beams:])[::-1]
+            )
+            top_token_indexes += current_top_token_indexes
+            top_cumulative_scores += list(
+                np.array(scores[step][beam_ix][current_top_token_indexes])
+                + current_beam.cumulative_score
+            )
+            beam_indexes += [beam_ix] * n_beams
+            current_completions += [beam_trees[beam_ix].current_sentence] * 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,
+                "token": top_tokens,
+            }
+        )
+        maxes = top_df.groupby(["token_index", "current_completions"])[
+            "cumulative_score"
+        ].idxmax()
+
+        top_df = top_df.loc[maxes]
 
-            if current_token_choice not in current_tree["children"]:
-                current_tree["children"][current_token_choice] = {
-                    "table": None,
-                    "children": {},
-                }
+        # Sort all top probabilities and keep top n_beams
+        top_df_selected = top_df.sort_values("cumulative_score", ascending=False).iloc[
+            :n_beams
+        ]
+        print(step)
+        display(top_df_selected)
 
-            # Rewrite the probs table even if it was there before, since new chosen nodes have appeared in the children of current tree
+        # 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
+        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]
+            print(step, beam_ix)
+            display(selected_tokens)
             markdown_table = generate_markdown_table(
-                step_scores[current_beam, :], current_sequence_score,
-                chosen_tokens=current_tree["children"].keys(),
+                step_scores[beam_ix, :],
+                current_beam.cumulative_score,
+                chosen_tokens=list(selected_tokens["token"].values),
             )
-            current_tree["table"] = markdown_table
+            beam_trees[beam_ix].table = markdown_table
 
-            current_tree = current_tree["children"][current_token_choice]
+        # Add new children for 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])
 
-            # Keep up to date the current cumulative score
-            current_sequence_score += step_scores[current_beam, current_token_choice_ix]
+            # Update the source tree
+            source_beam_ix = int(top_df_selected.iloc[beam_ix]["beam_index"])
 
-    display += generate_html(start_sentence, original_tree)
+            previous_len = len(str(original_tree))
+            beam_trees[source_beam_ix].children[current_token_choice] = 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(),
+            )
+            assert (
+                len(str(original_tree)) > previous_len
+            ), "Original tree has not increased size"
 
-    display += """
-        </ul>
-        </div>
-    </body>
-    """
-    return display
+        # Reassign all beams at once
+        beam_trees = [
+            beam_trees[int(top_df_selected.iloc[beam_ix]["beam_index"])]
+            for beam_ix in range(n_beams)
+        ]
+
+        # 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]
+    return original_tree
 
 @spaces.GPU
-def get_tables(input_text, number_steps, number_beams):
+def get_beam_search_html(input_text, number_steps, number_beams):
     inputs = tokenizer([input_text], return_tensors="pt")
 
     outputs = model.generate(
@@ -216,18 +319,18 @@ def get_tables(input_text, number_steps, number_beams):
         return_dict_in_generate=True,
         output_scores=True,
         top_k=5,
-        temperature=1.0,
-        do_sample=True,
+        do_sample=False,
     )
-    print(outputs.sequences_scores)
 
-    tables = display_tree(
+    original_tree = generate_beams(
         input_text,
-        outputs.scores,
-        outputs.sequences[:, len(inputs) :],
-        outputs.beam_indices[:, : -len(inputs)],
+        outputs.scores[:],
+        outputs.sequences[:, :],
+        outputs.beam_indices[:, :],
     )
-    return tables
+    html = generate_html(input_text, original_tree)
+    print(html)
+    return html
 
 
 with gr.Blocks(
@@ -241,6 +344,6 @@ with gr.Blocks(
     beams = gr.Slider(label="Number of beams", minimum=2, maximum=4, step=1, value=3)
     button = gr.Button()
     out = gr.Markdown(label="Output")
-    button.click(get_tables, inputs=[text, steps, beams], outputs=out)
+    button.click(get_beam_search_html, inputs=[text, steps, beams], outputs=out)
 
 demo.launch()