Update app.py

app.py

@@ -62,7 +62,7 @@ def ensure_float(value):
 def create_empty_figure(title):
     return go.Figure().update_layout(title=title, xaxis_title="", yaxis_title="", showlegend=False)
 
-# Function to process and visualize the full log probs with dynamic top_logprobs
+# Function to process and visualize the full log probs with dynamic top_logprobs, handling None
 def visualize_logprobs(json_input):
     try:
         # Parse the input (handles both JSON and Python dictionaries)
@@ -87,6 +87,9 @@ def visualize_logprobs(json_input):
                 logprobs.append(logprob)
                 # Get top_logprobs, default to empty dict if None
                 top_probs = entry.get("top_logprobs", {})
+                if top_probs is None:
+                    logger.debug("top_logprobs is None for token: %s, using empty dict", entry["token"])
+                    top_probs = {}  # Default to empty dict for None
                 # Ensure all values in top_logprobs are floats and create a list of tuples
                 finite_top_probs = []
                 for key, value in top_probs.items():
@@ -142,9 +145,12 @@ def visualize_logprobs(json_input):
         max_alternatives = max(len(alts) for alts in top_alternatives) if top_alternatives else 0
         for i, entry in enumerate(content):
             logprob = ensure_float(entry.get("logprob", None))
-            if math.isfinite(logprob) and logprob >= -100000 and "top_logprobs" in entry and entry["top_logprobs"] is not None:
+            if math.isfinite(logprob) and logprob >= -100000 and "top_logprobs" in entry:
                 token = entry["token"]
-                top_logprobs = entry["top_logprobs"]
+                top_logprobs = entry.get("top_logprobs", {})
+                if top_logprobs is None:
+                    logger.debug("top_logprobs is None for token: %s, using empty dict", token)
+                    top_logprobs = {}  # Default to empty dict for None
                 # Ensure all values in top_logprobs are floats
                 finite_top_logprobs = []
                 for key, value in top_logprobs.items():