Update app.py

app.py

@@ -5,12 +5,37 @@ import pandas as pd
 import io
 import base64
 import math
+import ast
+
+# Function to safely parse JSON or Python dictionary input
+def parse_input(json_input):
+    try:
+        # Try to parse as JSON first
+        data = json.loads(json_input)
+        return data
+    except json.JSONDecodeError as e:
+        try:
+            # If JSON fails, try to parse as Python literal (e.g., with single quotes)
+            data = ast.literal_eval(json_input)
+            # Convert Python dictionary to JSON-compatible format (replace single quotes with double quotes)
+            def dict_to_json(obj):
+                if isinstance(obj, dict):
+                    return {str(k): dict_to_json(v) for k, v in obj.items()}
+                elif isinstance(obj, list):
+                    return [dict_to_json(item) for item in obj]
+                else:
+                    return obj
+            return dict_to_json(data)
+        except (SyntaxError, ValueError) as e:
+            raise ValueError(f"Malformed input: {str(e)}. Ensure property names are in double quotes (e.g., \"content\") or correct Python dictionary format.")
 
 # Function to process and visualize log probs
 def visualize_logprobs(json_input):
     try:
-        # Parse the JSON input
-        data = json.loads(json_input)
+        # Parse the input (handles both JSON and Python dictionaries)
+        data = parse_input(json_input)
+        
+        # Ensure data is a list or dictionary with 'content'
         if isinstance(data, dict) and "content" in data:
             content = data["content"]
         elif isinstance(data, list):
@@ -130,13 +155,13 @@ def visualize_logprobs(json_input):
 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."
+        "Paste your JSON or Python dictionary log prob data below to visualize the tokens and their probabilities. Ensure property names are in double quotes (e.g., \"content\") for JSON, or use correct Python dictionary format."
     )
 
     json_input = gr.Textbox(
         label="JSON Input",
         lines=10,
-        placeholder="Paste your JSON here (ensure property names are in double quotes, e.g., \"token\")...",
+        placeholder="Paste your JSON (e.g., {\"content\": [...]}) or Python dict (e.g., {'content': [...]}) here...",
     )
 
     plot_output = gr.HTML(label="Log Probability Plot")