Upload app.py

app.py

@@ -39,8 +39,9 @@ def predict(prompt, dark_mode):
     normal_color = "white" if dark_mode else "black"
 
     for token, score in zip(first_beam_tokens, first_beam_scores):
-        # Replace SentencePiece space character (U+2581) and any underscores used for spaces
-        display_token = token.replace(' ', ' ').replace('_', ' ')
+        # Robustly replace SentencePiece space characters, literal underscores,
+        # and non-breaking spaces with regular spaces.
+        display_token = token.replace('▁', ' ').replace('_', ' ').replace(' ', ' ')
 
         if score < 0.6:
             color = "red"
@@ -61,7 +62,9 @@ def predict(prompt, dark_mode):
 
     detail_output = f'<table style="{detail_table_style}"><thead><tr><th style="{th_style}">Token</th><th style="{th_style}">Score</th></tr></thead><tbody>'
     for token, score in zip(first_beam_tokens, first_beam_scores):
-        display_token = token.replace(' ', ' ').replace('_', ' ')
+        # Robustly replace SentencePiece space characters, literal underscores,
+        # and non-breaking spaces with regular spaces.
+        display_token = token.replace(' ', ' ').replace('_', ' ').replace(' ', ' ')
 
         if score < 0.6:
             color = "red"
@@ -78,8 +81,12 @@ def predict(prompt, dark_mode):
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# LLM Token Probability Visualizer")
-    gr.Markdown("Input a prompt and see the first beam's tokens. Tokens with a probability score < 0.6 are highlighted in red.")
+    gr.Markdown("# InflectionLM: Beam Token Visualization")
+    gr.Markdown('''Input a prompt. The model will: 
+                - Generate multiple beams (a beam is a possible output generated by the model).
+                - Display the first output beam, highlighting any word or parts of words (tokens) with a probability score < 0.6 in red.
+                - Provide a toggle button to view detailed word/token probabilities for the beam.
+    ''')
 
     with gr.Column():
         prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=3)