Update app.py

app.py

@@ -4,12 +4,17 @@ import matplotlib.pyplot as plt
 import pandas as pd
 import io
 import base64
+import ast  # For safely evaluating Python literals
 
 # Function to process and visualize log probs
 def visualize_logprobs(json_input):
     try:
-        # Parse the JSON input
-        data = json.loads(json_input)
+        # Try to parse as JSON first
+        try:
+            data = json.loads(json_input)
+        except json.JSONDecodeError:
+            # If JSON fails, try to parse as Python literal (e.g., with single quotes)
+            data = ast.literal_eval(json_input)
         
         # Extract tokens and log probs, skipping None values
         tokens = [entry['token'] for entry in data['content'] if entry['logprob'] is not None]
@@ -67,10 +72,10 @@ def visualize_logprobs(json_input):
 # Gradio interface
 with gr.Blocks(title="Log Probability Visualizer") as app:
     gr.Markdown("# Log Probability Visualizer")
-    gr.Markdown("Paste your JSON log prob data below to visualize the tokens and their probabilities.")
+    gr.Markdown("Paste your JSON or Python dictionary log prob data below to visualize the tokens and their probabilities.")
     
     # Input
-    json_input = gr.Textbox(label="JSON Input", lines=10, placeholder="Paste your JSON here...")
+    json_input = gr.Textbox(label="JSON Input", lines=10, placeholder="Paste your JSON or Python dict here...")
     
     # Outputs
     plot_output = gr.HTML(label="Log Probability Plot")