Upload app.py

app.py

@@ -8,7 +8,7 @@ model, processor = start_model()
 print("Model loaded.")
 
 @spaces.GPU
-def predict(prompt):
+def predict(prompt, dark_mode):
     """
     Generates a response and returns an HTML string with tokens highlighted by probability.
     """
@@ -28,16 +28,24 @@ def predict(prompt):
     # Construct HTML output
     html_output = '<div style="font-family: monospace; white-space: pre-wrap; line-height: 1.5;">'
 
+    # Adjust colors based on dark mode
+    normal_color = "white" if dark_mode else "black"
+    border_color = "#444" if dark_mode else "#eee"
+
     for token, score in zip(first_beam_tokens, first_beam_scores):
-        # Replace BPE space marker with actual space for readability, but keep token boundaries
+        # Replace SentencePiece space character (U+2581) with a regular space
         display_token = token.replace(' ', ' ')
 
         if score < 0.6:
             color = "red"
+            # Light red background for low-probability tokens to make them stand out without breaking flow
+            bg_color = "#ffe6e6" if not dark_mode else "#441111"
+            style = f"color: {color}; background-color: {bg_color};"
         else:
-            color = "black"
+            color = normal_color
+            style = f"color: {color};"
 
-        html_output += f'<span style="color: {color}; border: 1px solid #eee; padding: 0 2px; margin: 0 1px;" title="Score: {score:.4f}">{display_token}</span>'
+        html_output += f'<span style="{style}" title="Score: {score:.4f}">{display_token}</span>'
 
     html_output += '</div>'
     return html_output
@@ -50,12 +58,13 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=3)
+            dark_mode_toggle = gr.Checkbox(label="Dark Mode", value=False)
             submit_btn = gr.Button("Generate")
 
         with gr.Column():
             output_html = gr.HTML(label="Highlighted Beam")
 
-    submit_btn.click(fn=predict, inputs=prompt_input, outputs=output_html)
+    submit_btn.click(fn=predict, inputs=[prompt_input, dark_mode_toggle], outputs=output_html)
 
 if __name__ == "__main__":
     demo.launch()