Update app.py

app.py

@@ -46,23 +46,24 @@ def parse_input(json_input):
 # Function to ensure a value is a float, converting from string if necessary
 def ensure_float(value):
     if value is None:
-        return None
+        logger.debug("Replacing None logprob with 0.0")
+        return 0.0  # Default to 0.0 for None to ensure visualization
     if isinstance(value, str):
         try:
             return float(value)
         except ValueError:
             logger.error("Failed to convert string '%s' to float", value)
-            return None
+            return 0.0  # Default to 0.0 for invalid strings
     if isinstance(value, (int, float)):
         return float(value)
-    return None
+    return 0.0  # Default for any other type
 
 # 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 log probs with interactive Plotly plots and pagination
-def visualize_logprobs(json_input, page_size=100, page=0):
+# Function to process and visualize the full log probs with dynamic top_logprobs
+def visualize_logprobs(json_input):
     try:
         # Parse the input (handles both JSON and Python dictionaries)
         data = parse_input(json_input)
@@ -75,47 +76,36 @@ def visualize_logprobs(json_input, page_size=100, page=0):
         else:
             raise ValueError("Input must be a list or dictionary with 'content' key")
 
-        # Extract tokens, log probs, and top alternatives, skipping None or non-finite values with fixed filter of -100000
+        # Extract tokens, log probs, and top alternatives, skipping non-finite values with fixed filter of -100000
         tokens = []
         logprobs = []
-        top_alternatives = []  # List to store top 3 log probs (selected token + 2 alternatives)
+        top_alternatives = []  # List to store all top_logprobs (dynamic length)
         for entry in content:
             logprob = ensure_float(entry.get("logprob", None))
-            if logprob is not None and math.isfinite(logprob) and logprob >= -100000:
+            if math.isfinite(logprob) and logprob >= -100000:
                 tokens.append(entry["token"])
                 logprobs.append(logprob)
                 # Get top_logprobs, default to empty dict if None
                 top_probs = entry.get("top_logprobs", {})
-                # Ensure all values in top_logprobs are floats
-                finite_top_probs = {}
+                # Ensure all values in top_logprobs are floats and create a list of tuples
+                finite_top_probs = []
                 for key, value in top_probs.items():
                     float_value = ensure_float(value)
                     if float_value is not None and math.isfinite(float_value):
-                        finite_top_probs[key] = float_value
-                # Get the top 3 log probs (including the selected token)
-                all_probs = {entry["token"]: logprob}  # Add the selected token's logprob
-                all_probs.update(finite_top_probs)  # Add alternatives
-                sorted_probs = sorted(all_probs.items(), key=lambda x: x[1], reverse=True)
-                top_3 = sorted_probs[:3]  # Top 3 log probs (highest to lowest)
-                top_alternatives.append(top_3)
+                        finite_top_probs.append((key, float_value))
+                # Sort by log probability (descending) to get all alternatives
+                sorted_probs = sorted(finite_top_probs, key=lambda x: x[1], reverse=True)
+                top_alternatives.append(sorted_probs)  # Store all alternatives, dynamic length
             else:
                 logger.debug("Skipping entry with logprob: %s (type: %s)", entry.get("logprob"), type(entry.get("logprob", None)))
 
         # Check if there's valid data after filtering
         if not logprobs or not tokens:
-            return (create_empty_figure("Log Probabilities of Generated Tokens"), None, "No finite log probabilities to display.", create_empty_figure("Top 3 Token Log Probabilities"), create_empty_figure("Significant Probability Drops"), 1, 0)
-
-        # Paginate data for large inputs
-        total_pages = max(1, (len(logprobs) + page_size - 1) // page_size)
-        start_idx = page * page_size
-        end_idx = min((page + 1) * page_size, len(logprobs))
-        paginated_tokens = tokens[start_idx:end_idx]
-        paginated_logprobs = logprobs[start_idx:end_idx]
-        paginated_alternatives = top_alternatives[start_idx:end_idx] if top_alternatives else []
+            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"))
 
         # 1. Main Log Probability Plot (Interactive Plotly)
         main_fig = go.Figure()
-        main_fig.add_trace(go.Scatter(x=list(range(len(paginated_logprobs))), y=paginated_logprobs, mode='markers+lines', name='Log Prob', marker=dict(color='blue')))
+        main_fig.add_trace(go.Scatter(x=list(range(len(logprobs))), y=logprobs, mode='markers+lines', name='Log Prob', marker=dict(color='blue')))
         main_fig.update_layout(
             title="Log Probabilities of Generated Tokens",
             xaxis_title="Token Position",
@@ -124,15 +114,15 @@ def visualize_logprobs(json_input, page_size=100, page=0):
             clickmode='event+select'
         )
         main_fig.update_traces(
-            customdata=[f"Token: {tok}, Log Prob: {prob:.4f}, Position: {i+start_idx}" for i, (tok, prob) in enumerate(zip(paginated_tokens, paginated_logprobs))],
+            customdata=[f"Token: {tok}, Log Prob: {prob:.4f}, Position: {i}" for i, (tok, prob) in enumerate(zip(tokens, logprobs))],
             hovertemplate='<b>%{customdata}</b><extra></extra>'
         )
 
         # 2. Probability Drop Analysis (Interactive Plotly)
-        if len(paginated_logprobs) < 2:
+        if len(logprobs) < 2:
             drops_fig = create_empty_figure("Significant Probability Drops")
         else:
-            drops = [paginated_logprobs[i+1] - paginated_logprobs[i] for i in range(len(paginated_logprobs)-1)]
+            drops = [logprobs[i+1] - logprobs[i] for i in range(len(logprobs)-1)]
             drops_fig = go.Figure()
             drops_fig.add_trace(go.Bar(x=list(range(len(drops))), y=drops, name='Drop', marker_color='red'))
             drops_fig.update_layout(
@@ -143,79 +133,75 @@ def visualize_logprobs(json_input, page_size=100, page=0):
                 clickmode='event+select'
             )
             drops_fig.update_traces(
-                customdata=[f"Drop: {drop:.4f}, From: {paginated_tokens[i]} to {paginated_tokens[i+1]}, Position: {i+start_idx}" for i, drop in enumerate(drops)],
+                customdata=[f"Drop: {drop:.4f}, From: {tokens[i]} to {tokens[i+1]}, Position: {i}" for i, drop in enumerate(drops)],
                 hovertemplate='<b>%{customdata}</b><extra></extra>'
             )
 
-        # Create DataFrame for the table (paginated)
+        # Create DataFrame for the table with dynamic top_logprobs
         table_data = []
-        for i, entry in enumerate(content[start_idx:end_idx]):
+        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 logprob is not None and 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 and entry["top_logprobs"] is not None:
                 token = entry["token"]
                 top_logprobs = entry["top_logprobs"]
                 # Ensure all values in top_logprobs are floats
-                finite_top_logprobs = {}
+                finite_top_logprobs = []
                 for key, value in top_logprobs.items():
                     float_value = ensure_float(value)
                     if float_value is not None and math.isfinite(float_value):
-                        finite_top_logprobs[key] = float_value
-                # Extract top 3 alternatives from top_logprobs
-                top_3 = sorted(finite_top_logprobs.items(), key=lambda x: x[1], reverse=True)[:3]
+                        finite_top_logprobs.append((key, float_value))
+                # Sort by log probability (descending)
+                sorted_probs = sorted(finite_top_logprobs, key=lambda x: x[1], reverse=True)
                 row = [token, f"{logprob:.4f}"]
-                for alt_token, alt_logprob in top_3:
+                for alt_token, alt_logprob in sorted_probs[:max_alternatives]:  # Use max number of alternatives
                     row.append(f"{alt_token}: {alt_logprob:.4f}")
-                while len(row) < 5:
+                # Pad with empty strings if fewer alternatives than max
+                while len(row) < 2 + max_alternatives:
                     row.append("")
                 table_data.append(row)
 
         df = (
             pd.DataFrame(
                 table_data,
-                columns=[
-                    "Token",
-                    "Log Prob",
-                    "Top 1 Alternative",
-                    "Top 2 Alternative",
-                    "Top 3 Alternative",
-                ],
+                columns=["Token", "Log Prob"] + [f"Alt {i+1}" for i in range(max_alternatives)],
             )
             if table_data
             else None
         )
 
-        # Generate colored text (paginated)
-        if paginated_logprobs:
-            min_logprob = min(paginated_logprobs)
-            max_logprob = max(paginated_logprobs)
+        # Generate colored text
+        if logprobs:
+            min_logprob = min(logprobs)
+            max_logprob = max(logprobs)
             if max_logprob == min_logprob:
-                normalized_probs = [0.5] * len(paginated_logprobs)
+                normalized_probs = [0.5] * len(logprobs)
             else:
                 normalized_probs = [
-                    (lp - min_logprob) / (max_logprob - min_logprob) for lp in paginated_logprobs
+                    (lp - min_logprob) / (max_logprob - min_logprob) for lp in logprobs
                 ]
 
             colored_text = ""
-            for i, (token, norm_prob) in enumerate(zip(paginated_tokens, normalized_probs)):
+            for i, (token, norm_prob) in enumerate(zip(tokens, normalized_probs)):
                 r = int(255 * (1 - norm_prob))  # Red for low confidence
                 g = int(255 * norm_prob)        # Green for high confidence
                 b = 0
                 color = f"rgb({r}, {g}, {b})"
                 colored_text += f'<span style="color: {color}; font-weight: bold;">{token}</span>'
-                if i < len(paginated_tokens) - 1:
+                if i < len(tokens) - 1:
                     colored_text += " "
             colored_text_html = f"<p>{colored_text}</p>"
         else:
             colored_text_html = "No finite log probabilities to display."
 
-        # Top 3 Token Log Probabilities (Interactive Plotly, paginated)
-        alt_viz_fig = create_empty_figure("Top 3 Token Log Probabilities") if not paginated_logprobs or not paginated_alternatives else go.Figure()
-        if paginated_logprobs and paginated_alternatives:
-            for i, (token, probs) in enumerate(zip(paginated_tokens, paginated_alternatives)):
+        # Top Token Log Probabilities (Interactive Plotly, dynamic length)
+        alt_viz_fig = create_empty_figure("Top Token Log Probabilities") if not logprobs or not top_alternatives else go.Figure()
+        if logprobs and top_alternatives:
+            for i, (token, probs) in enumerate(zip(tokens, top_alternatives)):
                 for j, (alt_tok, prob) in enumerate(probs):
-                    alt_viz_fig.add_trace(go.Bar(x=[f"{token} (Pos {i+start_idx})"], y=[prob], name=f"{alt_tok}", marker_color=['blue', 'green', 'red'][j]))
+                    alt_viz_fig.add_trace(go.Bar(x=[f"{token} (Pos {i})"], y=[prob], name=f"{alt_tok}", marker_color=['blue', 'green', 'red', 'purple', 'orange'][:len(probs)]))
             alt_viz_fig.update_layout(
-                title="Top 3 Token Log Probabilities (Paginated)",
+                title="Top Token Log Probabilities",
                 xaxis_title="Token (Position)",
                 yaxis_title="Log Probability",
                 barmode='stack',
@@ -223,33 +209,29 @@ def visualize_logprobs(json_input, page_size=100, page=0):
                 clickmode='event+select'
             )
             alt_viz_fig.update_traces(
-                customdata=[f"Token: {tok}, Alt: {alt}, Log Prob: {prob:.4f}, Position: {i+start_idx}" for i, (tok, alts) in enumerate(zip(paginated_tokens, paginated_alternatives)) for alt, prob in alts],
+                customdata=[f"Token: {tok}, Alt: {alt}, Log Prob: {prob:.4f}, Position: {i}" for i, (tok, alts) in enumerate(zip(tokens, top_alternatives)) for alt, prob in alts],
                 hovertemplate='<b>%{customdata}</b><extra></extra>'
             )
 
-        return (main_fig, df, colored_text_html, alt_viz_fig, drops_fig, total_pages, page)
+        return (main_fig, df, colored_text_html, alt_viz_fig, drops_fig)
 
     except Exception as e:
         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 3 Token Log Probabilities"), create_empty_figure("Significant Probability Drops"), 1, 0)
+        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 interactive layout and pagination
+# Gradio interface with full dataset visualization and dynamic top_logprobs
 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 the tokens and their probabilities. Fixed filter ≥ -100000, 100 tokens per page."
+        "Paste your JSON or Python dictionary log prob data below to visualize all tokens at once. Fixed filter ≥ -100000, dynamic number of top_logprobs."
     )
 
     with gr.Row():
-        with gr.Column(scale=1):
-            json_input = gr.Textbox(
-                label="JSON Input",
-                lines=10,
-                placeholder="Paste your JSON (e.g., {\"content\": [...]}) or Python dict (e.g., {'content': [...]}) here...",
-            )
-        with gr.Column(scale=1):
-            page = gr.Number(value=0, label="Page Number", precision=0, minimum=0)
-            page_size = gr.Number(value=100, label="Page Size", precision=0, minimum=10, maximum=1000, interactive=False)  # Fixed at 100, non-interactive
+        json_input = gr.Textbox(
+            label="JSON Input",
+            lines=10,
+            placeholder="Paste your JSON (e.g., {\"content\": [...]}) or Python dict (e.g., {'content': [...]}) here...",
+        )
 
     with gr.Row():
         plot_output = gr.Plot(label="Log Probability Plot (Click for Tokens)")
@@ -257,7 +239,7 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
 
     with gr.Row():
         table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
-        alt_viz_output = gr.Plot(label="Top 3 Token Log Probabilities (Click for Details)")
+        alt_viz_output = gr.Plot(label="Top Token Log Probabilities (Click for Details)")
 
     with gr.Row():
         text_output = gr.HTML(label="Colored Text (Confidence Visualization)")
@@ -265,46 +247,8 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
     btn = gr.Button("Visualize")
     btn.click(
         fn=visualize_logprobs,
-        inputs=[json_input, page_size, page],
-        outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, gr.State(), gr.State()],
-    )
-
-    # Pagination controls
-    with gr.Row():
-        prev_btn = gr.Button("Previous Page")
-        next_btn = gr.Button("Next Page")
-        total_pages_output = gr.Number(label="Total Pages", interactive=False)
-        current_page_output = gr.Number(label="Current Page", interactive=False)
-
-    def update_page(json_input, current_page, action):
-        try:
-            # Safely get total_pages by trying to process the data
-            result = visualize_logprobs(json_input, 100, 0)  # Use fixed page size and page 0
-            if isinstance(result[0], str) or result[0] is None:  # Check if it's an error message or empty figure
-                total_pages = 1  # Default to 1 page if no data
-            else:
-                total_pages = result[5]  # Extract total_pages from the result (index 5)
-        except Exception as e:
-            logger.error("Failed to calculate total pages: %s", str(e))
-            total_pages = 1  # Default to 1 page on error
-
-        if action == "prev" and current_page > 0:
-            current_page -= 1
-        elif action == "next":
-            if current_page < total_pages - 1:
-                current_page += 1
-        return gr.update(value=current_page), gr.update(value=total_pages)
-
-    prev_btn.click(
-        fn=update_page,
-        inputs=[json_input, page, gr.State()],
-        outputs=[page, total_pages_output]
-    )
-
-    next_btn.click(
-        fn=update_page,
-        inputs=[json_input, page, gr.State()],
-        outputs=[page, total_pages_output]
+        inputs=[json_input],
+        outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output],
     )
 
 app.launch()