Update app.py

app.py

@@ -371,9 +371,10 @@ def get_beam_search_html(input_text, number_steps, number_beams, length_penalty)
         output_scores=True,
         do_sample=False,
     )
-    print("Sequences:")
-    print(tokenizer.batch_decode(outputs.sequences))
-    print("Scores:", outputs.sequences_scores)
+    markdown = "Sequences:"
+    decoded_sequences = tokenizer.batch_decode(outputs.sequences)
+    for i, sequence in enumerate(decoded_sequences):
+        markdown += f"\n- {sequence} ( score {outputs.sequences_scores[i]:.2f})"
 
     original_tree = generate_beams(
         input_text,
@@ -382,7 +383,7 @@ def get_beam_search_html(input_text, number_steps, number_beams, length_penalty)
         length_penalty,
     )
     html = generate_html(input_text, original_tree)
-    return html
+    return html, markdown
 
 
 with gr.Blocks(
@@ -391,13 +392,24 @@ with gr.Blocks(
     ),
     css=STYLE,
 ) as demo:
+    gr.Markdown("""# Beam search visualizer
+                
+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.
+""")
     text = gr.Textbox(label="Sentence to decode from", value="Today is")
     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, length_penalty], outputs=out)
+    out_html = gr.Markdown()
+    out_markdown = gr.Markdown()
+    button.click(get_beam_search_html, inputs=[text, steps, beams, length_penalty], outputs=[out_html, out_markdown])
 
 demo.launch()