Update app.py

app.py

@@ -46,6 +46,7 @@ STYLE = """
 	height: auto;
 	text-align: center;
     display:inline-block;
+    padding-bottom: 10px!important;
 }
 #root {
     display: inline-grid!important;
@@ -417,6 +418,7 @@ def generate_beams(n_beams, start_sentence, scores, length_penalty, decoded_sequ
 
     return original_tree
 
+
 @spaces.GPU
 def get_beam_search_html(
     input_text, number_steps, number_beams, length_penalty, num_return_sequences
@@ -441,9 +443,12 @@ def get_beam_search_html(
     # Sequences are padded anyway so you can batch decode them
     decoded_sequences = tokenizer.batch_decode(outputs.sequences)
 
-    sequence_scores = (outputs.sequences_scores if number_beams > 1 else outputs.scores)
-    for i, sequence in enumerate(decoded_sequences):
-        markdown += f"\n- Score `{sequence_scores[i]:.2f}`: `{clean(sequence.replace('<s> ', ''))}`"
+    if number_beams > 1:
+        for i, sequence in enumerate(decoded_sequences):
+            markdown += f"\n- Score `{outputs.sequences_scores[i]:.2f}`: `{clean(sequence.replace('<s> ', ''))}`"
+    else:
+        markdown += f"\n- `{clean(decoded_sequences[0].replace('<s> ', ''))}`"
+    print(outputs.logits)
 
     if number_beams > 1:
         original_tree = generate_beams(
@@ -455,9 +460,9 @@ def get_beam_search_html(
         )
     else:
         original_tree = generate_beams(
-            n_beams,
+            number_beams,
             input_text,
-            outputs.logits,
+            outputs.scores,
             0,
             decoded_sequences,
         )