Update app.py

app.py

@@ -35,6 +35,7 @@ STYLE = """
     padding-right: 2px;
     padding-top: 0;
     padding-bottom: 0;
+    text-wrap:nowrap;
 }
 .tree {
 	padding: 0px;
@@ -137,10 +138,10 @@ a:before {
 	text-decoration-line: none;
 	border-radius: 5px;
 	transition: .5s;
-    width: 280px;
     display: flex;
     align-items: center;
-    justify-content: space-around;
+    justify-content: space-between;
+    overflow: hidden;
 }
 .tree li a span {
 	padding: 5px;
@@ -157,13 +158,19 @@ a:before {
 }
 .end-of-text, .chosen {
     background-color: #ea580c;
+}
+.end-of-text {
     width:auto!important;
 }
+.nonfinal {
+    width:280px;
+    min-width: 280px;
+}   
 """
 
 
 def clean(s):
-    return s.replace("\n", r"\n").replace("\t", r"\t")
+    return s.replace("\n", r"\n").replace("\t", r"\t").strip()
 
 
 def generate_markdown_table(
@@ -197,10 +204,10 @@ def generate_nodes(token_ix, node, step):
     token = tokenizer.decode([token_ix])
 
     if node.is_final:
-        return f"<li> <a href='#' class='end-of-text'> <span> <b>{token_ix}:<br>{clean(token)}</b> <br> Total score: {node.total_score:.2f} </span> </a> </li>"
+        return f"<li> <a href='#' class='end-of-text'> <span> <b>{clean(token)}</b> <br>Total score: {node.total_score:.2f}</span> </a> </li>"
 
     html_content = (
-        f"<li> <a href='#'> <span> <b>{token_ix}:<br>{clean(token)}</b> </span>"
+        f"<li> <a href='#' class='nonfinal'> <span> <b>{clean(token)}</b> </span>"
     )
     if node.table is not None:
         html_content += node.table
@@ -265,8 +272,6 @@ def generate_beams(start_sentence, scores, sequences, length_penalty):
     n_beams = len(scores[0])
     beam_trees = [original_tree] * n_beams
 
-    candidate_nodes = []
-
     for step, step_scores in enumerate(scores):
         (
             top_token_indexes,
@@ -359,16 +364,6 @@ def generate_beams(start_sentence, scores, sequences, length_penalty):
                 ),
             )
 
-            # Check this child should be selected as a top beam.
-            # Is it a final step or an EOS token?
-            if (
-                step == len(scores) - 1
-                or current_token_choice_ix == tokenizer.eos_token_id
-            ):
-                candidate_nodes.append(
-                    beam_trees[source_beam_ix].children[current_token_choice_ix]
-                )
-
         # Reassign all beams at once
         beam_trees = [
             beam_trees[int(top_df_selected.iloc[beam_ix]["beam_index"])]