Update app.py

app.py

@@ -61,8 +61,8 @@ 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 log probs with interactive Plotly plots
-def visualize_logprobs(json_input):
+# Function to process and visualize log probs with interactive Plotly plots and pagination
+def visualize_logprobs(json_input, prob_filter=-100000, page_size=100, page=0):
     try:
         # Parse the input (handles both JSON and Python dictionaries)
         data = parse_input(json_input)
@@ -75,13 +75,13 @@ def visualize_logprobs(json_input):
         else:
             raise ValueError("Input must be a list or dictionary with 'content' key")
 
-        # Extract tokens and log probs, skipping None or non-finite values with fixed filter of -100000
+        # Extract tokens, log probs, and top alternatives, skipping None or non-finite values with fixed filter
         tokens = []
         logprobs = []
         top_alternatives = []  # List to store top 3 log probs (selected token + 2 alternatives)
         for entry in content:
             logprob = ensure_float(entry.get("logprob", None))
-            if logprob is not None and math.isfinite(logprob) and logprob >= -100000:
+            if logprob is not None and math.isfinite(logprob) and logprob >= prob_filter:
                 tokens.append(entry["token"])
                 logprobs.append(logprob)
                 # Get top_logprobs, default to empty dict if None
@@ -103,11 +103,19 @@ def visualize_logprobs(json_input):
 
         # 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"))
+            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 (fixed page size of 100)
+        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 []
 
         # 1. Main Log Probability Plot (Interactive Plotly)
         main_fig = go.Figure()
-        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.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.update_layout(
             title="Log Probabilities of Generated Tokens",
             xaxis_title="Token Position",
@@ -116,15 +124,15 @@ def visualize_logprobs(json_input):
             clickmode='event+select'
         )
         main_fig.update_traces(
-            customdata=[f"Token: {tok}, Log Prob: {prob:.4f}, Position: {i}" for i, (tok, prob) in enumerate(zip(tokens, logprobs))],
+            customdata=[f"Token: {tok}, Log Prob: {prob:.4f}, Position: {i+start_idx}" for i, (tok, prob) in enumerate(zip(paginated_tokens, paginated_logprobs))],
             hovertemplate='<b>%{customdata}</b><extra></extra>'
         )
 
         # 2. Probability Drop Analysis (Interactive Plotly)
-        if len(logprobs) < 2:
+        if len(paginated_logprobs) < 2:
             drops_fig = create_empty_figure("Significant Probability Drops")
         else:
-            drops = [logprobs[i+1] - logprobs[i] for i in range(len(logprobs)-1)]
+            drops = [paginated_logprobs[i+1] - paginated_logprobs[i] for i in range(len(paginated_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(
@@ -135,15 +143,15 @@ def visualize_logprobs(json_input):
                 clickmode='event+select'
             )
             drops_fig.update_traces(
-                customdata=[f"Drop: {drop:.4f}, From: {tokens[i]} to {tokens[i+1]}, Position: {i}" for i, drop in enumerate(drops)],
+                customdata=[f"Drop: {drop:.4f}, From: {paginated_tokens[i]} to {paginated_tokens[i+1]}, Position: {i+start_idx}" for i, drop in enumerate(drops)],
                 hovertemplate='<b>%{customdata}</b><extra></extra>'
             )
 
-        # Create DataFrame for the table
+        # Create DataFrame for the table (paginated)
         table_data = []
-        for i, entry in enumerate(content):
+        for i, entry in enumerate(content[start_idx:end_idx]):
             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 logprob is not None and math.isfinite(logprob) and logprob >= prob_filter 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
@@ -176,38 +184,38 @@ def visualize_logprobs(json_input):
             else None
         )
 
-        # Generate colored text
-        if logprobs:
-            min_logprob = min(logprobs)
-            max_logprob = max(logprobs)
+        # Generate colored text (paginated)
+        if paginated_logprobs:
+            min_logprob = min(paginated_logprobs)
+            max_logprob = max(paginated_logprobs)
             if max_logprob == min_logprob:
-                normalized_probs = [0.5] * len(logprobs)
+                normalized_probs = [0.5] * len(paginated_logprobs)
             else:
                 normalized_probs = [
-                    (lp - min_logprob) / (max_logprob - min_logprob) for lp in logprobs
+                    (lp - min_logprob) / (max_logprob - min_logprob) for lp in paginated_logprobs
                 ]
 
             colored_text = ""
-            for i, (token, norm_prob) in enumerate(zip(tokens, normalized_probs)):
+            for i, (token, norm_prob) in enumerate(zip(paginated_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(tokens) - 1:
+                if i < len(paginated_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)
-        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)):
+        # 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)):
                 for j, (alt_tok, prob) in enumerate(probs):
-                    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.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.update_layout(
-                title="Top 3 Token Log Probabilities",
+                title="Top 3 Token Log Probabilities (Paginated)",
                 xaxis_title="Token (Position)",
                 yaxis_title="Log Probability",
                 barmode='stack',
@@ -215,29 +223,33 @@ def visualize_logprobs(json_input):
                 clickmode='event+select'
             )
             alt_viz_fig.update_traces(
-                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],
+                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],
                 hovertemplate='<b>%{customdata}</b><extra></extra>'
             )
 
-        return (main_fig, df, colored_text_html, alt_viz_fig, drops_fig)
+        return (main_fig, df, colored_text_html, alt_viz_fig, drops_fig, total_pages, page)
 
     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"))
+        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)
 
-# Gradio interface with improved layout
+# Gradio interface with interactive layout and pagination
 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, 1000 tokens per page."
+        "Paste your JSON or Python dictionary log prob data below to visualize the tokens and their probabilities. Fixed filter ≥ -100000, 100 tokens per page."
     )
 
     with gr.Row():
-        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):
+            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
 
     with gr.Row():
         plot_output = gr.Plot(label="Log Probability Plot (Click for Tokens)")
@@ -253,8 +265,36 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
     btn = gr.Button("Visualize")
     btn.click(
         fn=visualize_logprobs,
-        inputs=[json_input],
-        outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output],
+        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):
+        if action == "prev" and current_page > 0:
+            current_page -= 1
+        elif action == "next":
+            total_pages = visualize_logprobs(json_input, -100000, 100, 0)[5]  # Get total pages with fixed filter and page size
+            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]
     )
 
 app.launch()