Update app.py

app.py

@@ -58,11 +58,18 @@ def ensure_float(value):
         return float(value)
     return 0.0  # Default for any other type
 
+# Function to get or generate a token value (default to "Unknown" if missing)
+def get_token(entry):
+    token = entry.get("token", "Unknown")
+    if token == "Unknown":
+        logger.warning("Missing 'token' key for entry: %s, using 'Unknown'", entry)
+    return token
+
 # Function to create an empty Plotly figure
 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, handling None
+# Function to process and visualize the full log probs with dynamic top_logprobs, handling missing tokens
 def visualize_logprobs(json_input):
     try:
         # Parse the input (handles both JSON and Python dictionaries)
@@ -83,12 +90,13 @@ def visualize_logprobs(json_input):
         for entry in content:
             logprob = ensure_float(entry.get("logprob", None))
             if math.isfinite(logprob) and logprob >= -100000:
-                tokens.append(entry["token"])
+                token = get_token(entry)  # Safely get token, defaulting to "Unknown" if missing
+                tokens.append(token)
                 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"])
+                    logger.debug("top_logprobs is None for token: %s, using empty dict", 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 = []
@@ -146,7 +154,7 @@ def visualize_logprobs(json_input):
         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:
-                token = entry["token"]
+                token = get_token(entry)  # Safely get token, defaulting to "Unknown" if missing
                 top_logprobs = entry.get("top_logprobs", {})
                 if top_logprobs is None:
                     logger.debug("top_logprobs is None for token: %s, using empty dict", token)
@@ -225,11 +233,11 @@ def visualize_logprobs(json_input):
         logger.error("Visualization failed: %s", str(e))
         return (create_empty_figure("Log Probabilities of Generated Tokens"), None, "No finite log probabilities to display.", create_empty_figure("Top Token Log Probabilities"), create_empty_figure("Significant Probability Drops"))
 
-# Gradio interface with full dataset visualization and dynamic top_logprobs
+# Gradio interface with full dataset visualization, dynamic top_logprobs, and handling missing tokens
 with gr.Blocks(title="Log Probability Visualizer") as app:
     gr.Markdown("# Log Probability Visualizer")
     gr.Markdown(
-        "Paste your JSON or Python dictionary log prob data below to visualize all tokens at once. Fixed filter ≥ -100000, dynamic number of top_logprobs."
+        "Paste your JSON or Python dictionary log prob data below to visualize all tokens at once. Fixed filter ≥ -100000, dynamic number of top_logprobs, handles missing 'token'."
     )
 
     with gr.Row():