Update app.py

app.py

@@ -4,7 +4,7 @@ import matplotlib.pyplot as plt
 import pandas as pd
 import io
 import base64
-import ast  # For safely evaluating Python literals
+import ast
 
 # Function to process and visualize log probs
 def visualize_logprobs(json_input):
@@ -64,10 +64,34 @@ def visualize_logprobs(json_input):
             columns=["Token", "Log Prob", "Top 1 Alternative", "Top 2 Alternative", "Top 3 Alternative"]
         )
         
-        return img_html, df
+        # Generate colored text based on log probabilities
+        # Normalize log probs to [0, 1] for color scaling (0 = most uncertain, 1 = most confident)
+        min_logprob = min(logprobs)
+        max_logprob = max(logprobs)
+        if max_logprob == min_logprob:
+            normalized_probs = [0.5] * len(logprobs)  # Avoid division by zero
+        else:
+            normalized_probs = [(lp - min_logprob) / (max_logprob - min_logprob) for lp in logprobs]
+        
+        # Create HTML for colored text
+        colored_text = ""
+        for i, (token, norm_prob) in enumerate(zip(tokens, normalized_probs)):
+            # Map normalized probability to RGB color (green for high confidence, red for low)
+            # Use a simple linear interpolation: green (0, 255, 0) to red (255, 0, 0)
+            r = int(255 * (1 - norm_prob))  # Red increases as uncertainty increases
+            g = int(255 * norm_prob)        # Green decreases as uncertainty increases
+            b = 0                           # Blue stays 0 for simplicity
+            color = f'rgb({r}, {g}, {b})'
+            colored_text += f'<span style="color: {color}; font-weight: bold;">{token}</span>'
+            if i < len(tokens) - 1:
+                colored_text += " "  # Add space between tokens
+        
+        colored_text_html = f'<p>{colored_text}</p>'
+        
+        return img_html, df, colored_text_html
     
     except Exception as e:
-        return f"Error: {str(e)}", None
+        return f"Error: {str(e)}", None, None
 
 # Gradio interface
 with gr.Blocks(title="Log Probability Visualizer") as app:
@@ -80,13 +104,14 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
     # Outputs
     plot_output = gr.HTML(label="Log Probability Plot")
     table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
+    text_output = gr.HTML(label="Colored Text (Confidence Visualization)")
     
     # Button to trigger visualization
     btn = gr.Button("Visualize")
     btn.click(
         fn=visualize_logprobs,
         inputs=json_input,
-        outputs=[plot_output, table_output]
+        outputs=[plot_output, table_output, text_output]
     )
 
 # Launch the app