Update app.py

app.py

@@ -1,8 +1,8 @@
 import torch
-import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import numpy as np
 import gradio as gr
+import spaces
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 # True
@@ -122,6 +122,7 @@ 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."""
 
@@ -133,18 +134,17 @@ def generate_html(token, node):
         for token, subnode in node["children"].items():
             html_content += generate_html(token, subnode)
         html_content += "</ul>"
-
     html_content += "</li>"
-
     return html_content
 
 
-def generate_markdown_table(scores, top_k=4, chosen_tokens=None):
+def generate_markdown_table(scores, sequence_prob, top_k=4, chosen_tokens=None):
     markdown_table = """
     <table>
         <tr>
             <th><b>Token</b></th>
-            <th><b>Score</b></th>
+            <th><b>Step score</b></th>
+            <th><b>Cumulative score</b></th>
         </tr>"""
     for token_idx in np.argsort(scores)[-top_k:]:
         token = tokenizer.decode([token_idx])
@@ -155,6 +155,7 @@ def generate_markdown_table(scores, top_k=4, chosen_tokens=None):
         <tr style={style}>
             <td>{token}</td>
             <td>{scores[token_idx]:.4f}</td>
+            <td>{scores[token_idx] + sequence_prob:.4f}</td>
         </tr>"""
     markdown_table += """
     </table>"""
@@ -169,9 +170,11 @@ def display_tree(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 = tokenizer.decode([sequences[sequence_ix, step]])
+            current_token_choice_ix = sequences[sequence_ix, step]
+            current_token_choice = tokenizer.decode([current_token_choice_ix])
             current_beam = beam_indices[sequence_ix, step]
 
             if current_token_choice not in current_tree["children"]:
@@ -182,13 +185,16 @@ def display_tree(start_sentence, scores, sequences, beam_indices):
 
             # Rewrite the probs table even if it was there before, since new chosen nodes have appeared in the children of current tree
             markdown_table = generate_markdown_table(
-                step_scores[current_beam, :],
+                step_scores[current_beam, :], current_sequence_score,
                 chosen_tokens=current_tree["children"].keys(),
             )
             current_tree["table"] = markdown_table
 
             current_tree = current_tree["children"][current_token_choice]
 
+            # Keep up to date the current cumulative score
+            current_sequence_score += step_scores[current_beam, current_token_choice_ix]
+
     display += generate_html(start_sentence, original_tree)
 
     display += """
@@ -213,6 +219,7 @@ def get_tables(input_text, number_steps, number_beams):
         temperature=1.0,
         do_sample=True,
     )
+    print(outputs.sequences_scores)
 
     tables = display_tree(
         input_text,
@@ -221,7 +228,8 @@ def get_tables(input_text, number_steps, number_beams):
         outputs.beam_indices[:, : -len(inputs)],
     )
     return tables
-    
+
+
 with gr.Blocks(
     theme=gr.themes.Soft(
         text_size="lg", font=["monospace"], primary_hue=gr.themes.colors.green