Update app.py

app.py

@@ -7,15 +7,59 @@ import base64
 import ast
 import math
 
+# Function to safely convert string representations of infinity
+def parse_infinity(value):
+    if isinstance(value, str):
+        if value.lower() == '-infinity' or value.lower() == '-inf':
+            return float('-inf')
+        elif value.lower() == 'infinity' or value.lower() == 'inf':
+            return float('inf')
+    return value
+
 # Function to process and visualize log probs
 def visualize_logprobs(json_input):
     try:
-        # Try to parse as JSON first
+        # Try to parse as JSON first, handling string representations of infinity
         try:
-            data = json.loads(json_input)
+            # Attempt to load JSON, replacing -inf with "-Infinity" if needed
+            def replace_inf(s):
+                import re
+                return re.sub(r'-inf', '"-Infinity"', re.sub(r'inf', '"Infinity"', s))
+            
+            data = json.loads(replace_inf(json_input))
+            # Convert string "Infinity" or "-Infinity" back to float if needed
+            if isinstance(data, dict) and 'content' in data:
+                for entry in data['content']:
+                    if 'logprob' in entry:
+                        entry['logprob'] = parse_infinity(entry['logprob'])
+                    if 'top_logprobs' in entry:
+                        entry['top_logprobs'] = {k: parse_infinity(v) for k, v in entry['top_logprobs'].items()}
+            elif isinstance(data, list):
+                for entry in data:
+                    if 'logprob' in entry:
+                        entry['logprob'] = parse_infinity(entry['logprob'])
+                    if 'top_logprobs' in entry:
+                        entry['top_logprobs'] = {k: parse_infinity(v) for k, v in entry['top_logprobs'].items()}
+            
         except json.JSONDecodeError:
             # If JSON fails, try to parse as Python literal (e.g., with single quotes)
-            data = ast.literal_eval(json_input)
+            try:
+                data = ast.literal_eval(json_input)
+                # Ensure -inf is handled as float('-inf')
+                if isinstance(data, dict) and 'content' in data:
+                    for entry in data['content']:
+                        if 'logprob' in entry and isinstance(entry['logprob'], str):
+                            entry['logprob'] = parse_infinity(entry['logprob'])
+                        if 'top_logprobs' in entry:
+                            entry['top_logprobs'] = {k: parse_infinity(v) for k, v in entry['top_logprobs'].items()}
+                elif isinstance(data, list):
+                    for entry in data:
+                        if 'logprob' in entry and isinstance(entry['logprob'], str):
+                            entry['logprob'] = parse_infinity(entry['logprob'])
+                        if 'top_logprobs' in entry:
+                            entry['top_logprobs'] = {k: parse_infinity(v) for k, v in entry['top_logprobs'].items()}
+            except (SyntaxError, ValueError) as e:
+                raise ValueError(f"Malformed input: {str(e)}")
         
         # Ensure data is a list or dictionary with 'content'
         if isinstance(data, dict) and 'content' in data:
@@ -29,17 +73,17 @@ def visualize_logprobs(json_input):
         tokens = []
         logprobs = []
         for entry in content:
-            if entry['logprob'] is not None and math.isfinite(entry['logprob']):
+            if 'logprob' in entry and entry['logprob'] is not None and math.isfinite(entry['logprob']):
                 tokens.append(entry['token'])
                 logprobs.append(entry['logprob'])
         
         # Prepare data for the table
         table_data = []
         for entry in content:
-            if entry['logprob'] is not None and math.isfinite(entry['logprob']):
+            if 'logprob' in entry and entry['logprob'] is not None and math.isfinite(entry['logprob']):
                 token = entry['token']
                 logprob = entry['logprob']
-                top_logprobs = entry['top_logprobs'] or {}
+                top_logprobs = entry.get('top_logprobs', {})
                 
                 # Filter out non-finite (e.g., -inf, inf, nan) log probs from top_logprobs
                 finite_top_logprobs = {k: v for k, v in top_logprobs.items() if math.isfinite(v)}