Update app.py

app.py

@@ -281,7 +281,7 @@ def visualize_logprobs(json_input, chunk=0, chunk_size=100):
         logger.error("Visualization failed: %s (Input: %s)", str(e), json_input[:100] + "..." if len(json_input) > 100 else json_input)
         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, 0)
 
-# Analysis functions for detecting correct vs. incorrect traces (unchanged from previous)
+# Analysis functions for detecting correct vs. incorrect traces (unchanged)
 def analyze_confidence_signature(logprobs, tokens):
     if not logprobs or not tokens:
         return "No data for confidence signature analysis.", None
@@ -476,114 +476,120 @@ def analyze_full_trace(json_input):
         return analysis_html, None, None, None, None, None
 
 # Gradio interface with two tabs: Trace Analysis and Visualization
-with gr.Blocks(title="Log Probability Visualizer") as app:
-    gr.Markdown("# Log Probability Visualizer")
-    gr.Markdown(
-        "Paste your JSON log prob data below to analyze reasoning traces and visualize tokens in chunks of 100. Fixed filter ≥ -100000, dynamic number of top_logprobs, handles missing or null fields. Next chunk is precomputed proactively."
-    )
+try:
+    with gr.Blocks(title="Log Probability Visualizer") as app:
+        gr.Markdown("# Log Probability Visualizer")
+        gr.Markdown(
+            "Paste your JSON log prob data below to analyze reasoning traces and visualize tokens in chunks of 100. Fixed filter ≥ -100000, dynamic number of top_logprobs, handles missing or null fields. Next chunk is precomputed proactively."
+        )
 
-    with gr.Tabs():
-        with gr.Tab("Trace Analysis"):
-            with gr.Row():
-                json_input_analysis = gr.Textbox(
-                    label="JSON Input for Trace Analysis",
-                    lines=10,
-                    placeholder="Paste your JSON (e.g., {\"content\": [{\"bytes\": [44], \"logprob\": 0.0, \"token\": \",\", \"top_logprobs\": {\" so\": -13.8046875, \".\": -13.8046875, \"，\": -13.640625}}]}).",
+        with gr.Tabs():
+            with gr.Tab("Trace Analysis"):
+                with gr.Row():
+                    json_input_analysis = gr.Textbox(
+                        label="JSON Input for Trace Analysis",
+                        lines=10,
+                        placeholder="Paste your JSON (e.g., {\"content\": [{\"bytes\": [44], \"logprob\": 0.0, \"token\": \",\", \"top_logprobs\": {\" so\": -13.8046875, \".\": -13.8046875, \"，\": -13.640625}}]}).",
+                    )
+                with gr.Row():
+                    analysis_output = gr.HTML(label="Trace Analysis Results")
+
+                btn_analyze = gr.Button("Analyze Trace")
+                btn_analyze.click(
+                    fn=analyze_full_trace,
+                    inputs=[json_input_analysis],
+                    outputs=[analysis_output, gr.State(), gr.State(), gr.State(), gr.State(), gr.State()],
                 )
-            with gr.Row():
-                analysis_output = gr.HTML(label="Trace Analysis Results")
-
-            btn_analyze = gr.Button("Analyze Trace")
-            btn_analyze.click(
-                fn=analyze_full_trace,
-                inputs=[json_input_analysis],
-                outputs=[analysis_output, gr.State(), gr.State(), gr.State(), gr.State(), gr.State()],
-            )
 
-        with gr.Tab("Visualization"):
-            with gr.Row():
-                json_input_viz = gr.Textbox(
-                    label="JSON Input for Visualization",
-                    lines=10,
-                    placeholder="Paste your JSON (e.g., {\"content\": [{\"bytes\": [44], \"logprob\": 0.0, \"token\": \",\", \"top_logprobs\": {\" so\": -13.8046875, \".\": -13.8046875, \"，\": -13.640625}}]}).",
+            with gr.Tab("Visualization"):
+                with gr.Row():
+                    json_input_viz = gr.Textbox(
+                        label="JSON Input for Visualization",
+                        lines=10,
+                        placeholder="Paste your JSON (e.g., {\"content\": [{\"bytes\": [44], \"logprob\": 0.0, \"token\": \",\", \"top_logprobs\": {\" so\": -13.8046875, \".\": -13.8046875, \"，\": -13.640625}}]}).",
+                    )
+                    chunk = gr.Number(value=0, label="Current Chunk", precision=0, minimum=0)
+
+                with gr.Row():
+                    plot_output = gr.Plot(label="Log Probability Plot (Click for Tokens)")
+                    drops_output = gr.Plot(label="Probability Drops (Click for Details)")
+
+                with gr.Row():
+                    table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
+                    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)")
+
+                with gr.Row():
+                    prev_btn = gr.Button("Previous Chunk")
+                    next_btn = gr.Button("Next Chunk")
+                    total_chunks_output = gr.Number(label="Total Chunks", interactive=False)
+
+                # Precomputed next chunk state (hidden)
+                precomputed_next = gr.State(value=None)
+
+                btn_viz = gr.Button("Visualize")
+                btn_viz.click(
+                    fn=visualize_logprobs,
+                    inputs=[json_input_viz, chunk],
+                    outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
                 )
-                chunk = gr.Number(value=0, label="Current Chunk", precision=0, minimum=0)
-
-            with gr.Row():
-                plot_output = gr.Plot(label="Log Probability Plot (Click for Tokens)")
-                drops_output = gr.Plot(label="Probability Drops (Click for Details)")
-
-            with gr.Row():
-                table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
-                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)")
-
-            with gr.Row():
-                prev_btn = gr.Button("Previous Chunk")
-                next_btn = gr.Button("Next Chunk")
-                total_chunks_output = gr.Number(label="Total Chunks", interactive=False)
-
-            # Precomputed next chunk state (hidden)
-            precomputed_next = gr.State(value=None)
-
-            btn_viz = gr.Button("Visualize")
-            btn_viz.click(
-                fn=visualize_logprobs,
-                inputs=[json_input_viz, chunk],
-                outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
-            )
-
-            # Precompute next chunk proactively when on current chunk
-            async def precompute_next_chunk(json_input, current_chunk, precomputed_next):
-                if precomputed_next is not None:
-                    return precomputed_next  # Use cached precomputed chunk if available
-                next_tokens, next_logprobs, next_alternatives = await precompute_chunk(json_input, 100, current_chunk)
-                if next_tokens is None or next_logprobs is None or next_alternatives is None:
-                    return None
-                return (next_tokens, next_logprobs, next_alternatives)
-
-            # Update chunk on button clicks
-            def update_chunk(json_input, current_chunk, action, precomputed_next=None):
-                total_chunks = visualize_logprobs(json_input, 0)[5]  # Get total chunks
-                if action == "prev" and current_chunk > 0:
-                    current_chunk -= 1
-                elif action == "next" and current_chunk < total_chunks - 1:
-                    current_chunk += 1
-                    # If precomputed next chunk exists, use it; otherwise, compute it
-                    if precomputed_next:
-                        next_tokens, next_logprobs, next_alternatives = precomputed_next
-                        if next_tokens and next_logprobs and next_alternatives:
-                            logger.debug("Using precomputed next chunk for chunk %d", current_chunk)
-                            return visualize_logprobs(json_input, current_chunk)
-                return visualize_logprobs(json_input, current_chunk)
-
-            prev_btn.click(
-                fn=update_chunk,
-                inputs=[json_input_viz, chunk, gr.State(value="prev"), precomputed_next],
-                outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
-            )
-
-            next_btn.click(
-                fn=update_chunk,
-                inputs=[json_input_viz, chunk, gr.State(value="next"), precomputed_next],
-                outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
-            )
+                # Precompute next chunk proactively when on current chunk
+                async def precompute_next_chunk(json_input, current_chunk, precomputed_next):
+                    if precomputed_next is not None:
+                        return precomputed_next  # Use cached precomputed chunk if available
+                    try:
+                        next_tokens, next_logprobs, next_alternatives = await precompute_chunk(json_input, 100, current_chunk)
+                        if next_tokens is None or next_logprobs is None or next_alternatives is None:
+                            return None
+                        return (next_tokens, next_logprobs, next_alternatives)
+                    except Exception as e:
+                        logger.error("Precomputation failed for chunk %d: %s", current_chunk + 1, str(e))
+                        return None
+
+                # Update chunk on button clicks
+                def update_chunk(json_input, current_chunk, action, precomputed_next=None):
+                    total_chunks = visualize_logprobs(json_input, 0)[5]  # Get total chunks
+                    if action == "prev" and current_chunk > 0:
+                        current_chunk -= 1
+                    elif action == "next" and current_chunk < total_chunks - 1:
+                        current_chunk += 1
+                        # If precomputed next chunk exists, use it; otherwise, compute it
+                        if precomputed_next:
+                            next_tokens, next_logprobs, next_alternatives = precomputed_next
+                            if next_tokens and next_logprobs and next_alternatives:
+                                logger.debug("Using precomputed next chunk for chunk %d", current_chunk)
+                                return visualize_logprobs(json_input, current_chunk)
+                    return visualize_logprobs(json_input, current_chunk)
+
+                prev_btn.click(
+                    fn=update_chunk,
+                    inputs=[json_input_viz, chunk, gr.State(value="prev"), precomputed_next],
+                    outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
+                )
 
-            # Trigger precomputation when chunk changes (via button clicks or initial load)
-            def trigger_precomputation(json_input, current_chunk):
-                try:
-                    asyncio.create_task(precompute_next_chunk(json_input, current_chunk, None))
-                except Exception as e:
-                    logger.error("Precomputation trigger failed: %s", str(e))
-                return gr.update(value=current_chunk)
-
-            # Use a dummy event to trigger precomputation on chunk change (simplified for Gradio)
-            chunk.change(
-                fn=trigger_precomputation,
-                inputs=[json_input_viz, chunk],
-                outputs=[chunk],
-            )
+                next_btn.click(
+                    fn=update_chunk,
+                    inputs=[json_input_viz, chunk, gr.State(value="next"), precomputed_next],
+                    outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
+                )
 
-app.launch()
+                # Trigger precomputation when chunk changes (via button clicks or initial load)
+                def trigger_precomputation(json_input, current_chunk):
+                    try:
+                        asyncio.create_task(precompute_next_chunk(json_input, current_chunk, None))
+                    except Exception as e:
+                        logger.error("Precomputation trigger failed: %s", str(e))
+                    return gr.update(value=current_chunk)
+
+                # Use a dummy event to trigger precomputation on chunk change (simplified for Gradio)
+                chunk.change(
+                    fn=trigger_precomputation,
+                    inputs=[json_input_viz, chunk],
+                    outputs=[chunk],
+                )
+except Exception as e:
+    logger.error("Application startup failed: %s", str(e))
+    raise