Update app.py

app.py

@@ -62,7 +62,7 @@ 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, prob_filter=-100000, page_size=100, page=0):
+def visualize_logprobs(json_input, 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, prob_filter=-100000, 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
+        # Extract tokens, log probs, and top alternatives, 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)
         for entry in content:
             logprob = ensure_float(entry.get("logprob", None))
-            if logprob is not None and math.isfinite(logprob) and logprob >= prob_filter:
+            if logprob is not None and math.isfinite(logprob) and logprob >= -100000:
                 tokens.append(entry["token"])
                 logprobs.append(logprob)
                 # Get top_logprobs, default to empty dict if None
@@ -105,7 +105,7 @@ def visualize_logprobs(json_input, prob_filter=-100000, page_size=100, page=0):
         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 (fixed page size of 100)
+        # 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))
@@ -151,7 +151,7 @@ def visualize_logprobs(json_input, prob_filter=-100000, page_size=100, page=0):
         table_data = []
         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 >= prob_filter and "top_logprobs" in entry and entry["top_logprobs"] is not 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"]
                 top_logprobs = entry["top_logprobs"]
                 # Ensure all values in top_logprobs are floats
@@ -277,10 +277,20 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
         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":
-            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)