Update app.py

app.py

@@ -24,15 +24,16 @@ STYLE = """
     overflow: scroll;
 }
 .prose ul ul {
-    margin: 0!important;
     font-size: 10px!important;
 }
+.prose li {
+    margin-bottom: 0!important;
+}
 .prose td, th {
     padding-left: 2px;
     padding-right: 2px;
     padding-top: 0;
     padding-bottom: 0;
-    text-wrap: nowrap;
 }
 
 .tree {
@@ -42,69 +43,77 @@ STYLE = """
     font-size: 10px;
 	width: 100%;
 	height: auto;
+	text-align: center;
+    display:inline-block;
 }
 .tree ul {
-	padding-top: 20px;
-	position: relative;
-	transition: .5s;
-    margin: 0!important;
+    padding-left: 20px;
+    position: relative;
+    transition: all 0.5s ease 0s;
     display: flex;
-    flex-direction: row;
-    justify-content: center;
-    gap:10px;
+    flex-direction: column;
+    gap: 10px;
+    margin: 0px !important;
 }
 .tree li {
-	display: inline-table;
-	text-align: center;
-	list-style-type: none;
-	position: relative;
-	padding-top: 10px;
-	transition: .5s;
+    display: flex;
+    text-align: center;
+    list-style-type: none;
+    position: relative;
+    padding-left: 20px;
+    transition: all 0.5s ease 0s;
+    flex-direction: row;
+    justify-content: start;
+    align-items: center;
 }
+
 .tree li::before, .tree li::after {
-	content: '';
-	position: absolute;
-	top: 0;
-	right: 50%;
-	border-top: 2px solid var(--body-text-color);
-	width: 55%;
-    min-width: 30px;
-	height: 10px;
+    content: "";
+    position: absolute;
+    left: 0px;
+    border-left: 1px solid var(--body-text-color);
+    width: 20px;
 }
-.tree li::after {
-	right: auto;
-	left: 50%;
-	border-left: 2px solid var(--body-text-color);
+.tree li::before {
+    top: 0;
+    height:50%;
 }
-
-.tree li:only-child::after, .tree li:only-child::before {
-	display: none;
+.tree li::after {
+    top: 50%;
+    height: 70%;
+    bottom: auto;
+    border-top: 1px solid var(--body-text-color);
 }
-.tree ul:has(> li:only-child)::before {
-    height:40px;
+.tree li:only-child::after, li:only-child::before {
+    display: none;
 }
+
 .tree li:first-child::before, .tree li:last-child::after {
-	border: 0 none;
+    border: 0 none;
 }
 .tree li:last-child::before {
-	border-right: 2px solid var(--body-text-color);
-	border-radius: 0 5px 0 0;
-	-webkit-border-radius: 0 5px 0 0;
-	-moz-border-radius: 0 5px 0 0;
+	border-bottom: 1px solid var(--body-text-color);
+	border-radius: 0px 0px 0px 5px;
+	-webkit-border-radius: 0px 0px 0px 5px;
+	-moz-border-radius: 0px 0px 0px 5px;
 }
 .tree li:first-child::after {
 	border-radius: 5px 0 0 0;
 	-webkit-border-radius: 5px 0 0 0;
 	-moz-border-radius: 5px 0 0 0;
 }
+
 .tree ul ul::before {
-	content: '';
-	position: absolute;
-	top: 0;
-	left: 50%;
-	border-left: 2px solid var(--body-text-color);
-	width: 0;
-	height: 20px;
+    content: "";
+    position: absolute;
+    left: 0;
+    top: 50%;
+    border-top: 1px solid var(--body-text-color);
+    width: 20px;
+    height: 0;
+}
+.tree ul:has(> li:only-child)::before {
+    width:40px;
 }
 .tree li a {
 	border: 1px solid var(--body-text-color);
@@ -114,6 +123,7 @@ STYLE = """
 	text-decoration-line: none;
 	border-radius: 5px;
 	transition: .5s;
+    width: 200px;
 }
 .tree li a span {
 	padding: 5px;
@@ -127,8 +137,7 @@ STYLE = """
 	background: #ffedd5;
 }
 .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: #f97316;
-    color= #CCC;
+	border-color: #7c2d12;
 }
 .chosen {
     background-color: #ea580c;
@@ -176,6 +185,7 @@ def generate_markdown_table(scores, sequence_prob, top_k=4, chosen_tokens=None):
 
 
 def generate_html(start_sentence, original_tree):
+
     html_output = """<div class="custom-container">
 				<div class="tree">
                 <ul>"""
@@ -203,6 +213,7 @@ class BeamNode:
 
 
 def generate_beams(start_sentence, scores, sequences, beam_indices):
+    print(tokenizer.batch_decode(sequences))
     sequences = sequences.cpu().numpy()
     original_tree = BeamNode(
         cumulative_score=0, table=None, current_sentence=start_sentence, children={}
@@ -230,9 +241,7 @@ def generate_beams(start_sentence, scores, sequences, beam_indices):
             )
             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_tokens += [tokenizer.decode([el]) for el in current_top_token_indexes]
 
         top_df = pd.DataFrame.from_dict(
             {
@@ -258,7 +267,9 @@ def generate_beams(start_sentence, scores, sequences, beam_indices):
         # 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]
+            selected_tokens = top_df_selected.loc[
+                top_df_selected["beam_index"] == beam_ix
+            ]
             markdown_table = generate_markdown_table(
                 step_scores[beam_ix, :],
                 current_beam.cumulative_score,
@@ -323,6 +334,7 @@ def get_beam_search_html(input_text, number_steps, number_beams):
         outputs.beam_indices[:, :],
     )
     html = generate_html(input_text, original_tree)
+    print(html)
     return html
 
 
@@ -333,7 +345,7 @@ 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=7, step=1, value=4)
+    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)
     button = gr.Button()
     out = gr.Markdown(label="Output")