Upload app.py

app.py

@@ -28,16 +28,20 @@ def predict(prompt, dark_mode):
     first_beam_score = exp(generated_dicts.sequences_scores[0]).item()
 
     # Construct Main HTML output
-    # We use a <h3> tag to simulate a Markdown header for the beam title and score
-    html_output = f'<h3 style="font-family: sans-serif;">Beam 1 | Score: {first_beam_score:.4f}</h3>'
-    html_output += '<div style="font-family: monospace; white-space: pre-wrap; line-height: 1.5;">'
+    # We use a simple div for the title instead of a markdown-style header
+    main_container_style = "border: 2px solid lightblue; padding: 15px; border-radius: 8px; font-family: sans-serif;"
+    title_style = "font-weight: bold; margin-bottom: 10px; font-size: 1.1em;"
+
+    html_output = f'<div style="{main_container_style}">'
+    html_output += f'<div style="{title_style}">Beam 1 | Score: {first_beam_score:.4f}</div>'
+    html_output += '<div style="font-family: monospace; white-space: pre-wrap; line-height: 1.5; overflow-wrap: break-word; word-wrap: break-word;">'
 
     # Adjust colors based on 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) with a regular space
-        display_token = token.replace(' ', ' ')
+        # Replace SentencePiece space character (U+2581) and any underscores used for spaces
+        display_token = token.replace(' ', ' ').replace('_', ' ')
 
         if score < 0.6:
             color = "red"
@@ -49,7 +53,7 @@ def predict(prompt, dark_mode):
 
         html_output += f'<span style="{style}" title="Score: {score:.4f}">{display_token}</span>'
 
-    html_output += '</div>'
+    html_output += '</div></div>'
 
     # Construct Detailed HTML output
     # We create a simple table to list tokens and their scores
@@ -59,8 +63,19 @@ 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(' ', ' ')
-        detail_output += f'<tr><td style="{td_style}">{display_token}</td><td style="{td_style}">{score:.4f}</td></tr>'
+        # Replace SentencePiece space character (U+2581) and any underscores used for spaces
+        display_token = token.replace(' ', ' ').replace('_', ' ')
+
+        # Apply red highlighting for low-probability tokens
+        if score < 0.6:
+            color = "red"
+            bg_color = "#441111" if dark_mode else "#ffe6e6"
+            token_cell_style = f"{td_style} color: {color}; background-color: {bg_color};"
+        else:
+            color = normal_color
+            token_cell_style = f"{td_style} color: {color};"
+
+        detail_output += f'<tr><td style="{token_cell_style}">{display_token}</td><td style="{td_style}">{score:.4f}</td></tr>'
     detail_output += '</tbody></table>'
 
     return html_output, detail_output