Update app.py

app.py

@@ -9,7 +9,6 @@ import ast
 import logging
 import numpy as np
 import plotly.graph_objects as go
-from plotly.subplots import make_subplots
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
@@ -58,8 +57,12 @@ def ensure_float(value):
         return float(value)
     return None
 
+# 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
-def visualize_logprobs(json_input, page=0):
+def visualize_logprobs(json_input):
     try:
         # Parse the input (handles both JSON and Python dictionaries)
         data = parse_input(json_input)
@@ -72,7 +75,7 @@ def visualize_logprobs(json_input, 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 and log probs, skipping None or non-finite values with fixed filter of -100000
         tokens = []
         logprobs = []
         top_alternatives = []  # List to store top 3 log probs (selected token + 2 alternatives)
@@ -100,20 +103,11 @@ def visualize_logprobs(json_input, page=0):
 
         # Check if there's valid data after filtering
         if not logprobs or not tokens:
-            return (gr.update(value="No finite log probabilities or tokens to visualize after filtering"), None, None, None, 1, 0)
-
-        # Paginate data for large inputs (fixed page size of 1000)
-        page_size = 1000
-        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 3 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",
@@ -122,16 +116,15 @@ def visualize_logprobs(json_input, 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:
-            drops_fig = go.Figure()
-            drops_fig.add_trace(go.Bar(x=list(range(len(paginated_logprobs)-1)), y=[0], name='Drop', marker_color='red'))
+        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(
@@ -142,13 +135,13 @@ def visualize_logprobs(json_input, 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
         table_data = []
-        for i, entry in enumerate(content[start_idx:end_idx]):
+        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:
                 token = entry["token"]
@@ -183,38 +176,38 @@ def visualize_logprobs(json_input, page=0):
             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 (paginated)
-        alt_viz_fig = go.Figure()
-        if paginated_logprobs and paginated_alternatives:
-            for i, (token, probs) in enumerate(zip(paginated_tokens, paginated_alternatives)):
+        # Top 3 Token Log Probabilities (Interactive Plotly)
+        alt_viz_fig = create_empty_figure("Top 3 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'][j]))
             alt_viz_fig.update_layout(
-                title="Top 3 Token Log Probabilities (Paginated)",
+                title="Top 3 Token Log Probabilities",
                 xaxis_title="Token (Position)",
                 yaxis_title="Log Probability",
                 barmode='stack',
@@ -222,35 +215,29 @@ def visualize_logprobs(json_input, 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>'
             )
-            alt_viz_html = alt_viz_fig.to_html(include_plotlyjs='cdn', full_html=False)
-        else:
-            alt_viz_html = "No finite log probabilities to display."
 
-        return (main_fig, df, colored_text_html, alt_viz_html, 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 (gr.update(value=f"Error: {str(e)}"), None, "No finite log probabilities to display.", None, gr.update(value="No data for probability drops."), 1, 0)
+        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"))
 
-# Gradio interface with interactive layout and pagination
+# Gradio interface with improved layout
 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. Use pagination to navigate large inputs (fixed filter ≥ -100000, 1000 tokens per page)."
+        "Paste your JSON or Python dictionary log prob data below to visualize the tokens and their probabilities. Fixed filter ≥ -100000, 1000 tokens per page."
     )
 
     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)
+        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)")
@@ -266,36 +253,8 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
     btn = gr.Button("Visualize")
     btn.click(
         fn=visualize_logprobs,
-        inputs=[json_input, 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):
-        if action == "prev" and current_page > 0:
-            current_page -= 1
-        elif action == "next":
-            total_pages = visualize_logprobs(json_input, 0)[5]  # Get total pages
-            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()