Update app.py

app.py

@@ -8,8 +8,6 @@ import math
 import logging
 import numpy as np
 import plotly.graph_objects as go
-import asyncio
-import threading
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
@@ -30,7 +28,7 @@ def parse_input(json_input):
 def ensure_float(value):
     if value is None:
         logger.debug("Replacing None logprob with 0.0")
-        return 0.0  # Default to 0.0 for None to ensure visualization
+        return 0.0  # Default to 0.0 for None
     if isinstance(value, str):
         try:
             return float(value)
@@ -52,8 +50,8 @@ def get_token(entry):
 def create_empty_figure(title):
     return go.Figure().update_layout(title=title, xaxis_title="", yaxis_title="", showlegend=False)
 
-# Precompute the next chunk asynchronously
-async def precompute_chunk(json_input, chunk_size, current_chunk):
+# Precompute the next chunk (synchronous for Hugging Face Spaces)
+def precompute_chunk(json_input, chunk_size, current_chunk):
     try:
         data = parse_input(json_input)
         content = data.get("content", []) if isinstance(data, dict) else data
@@ -71,10 +69,7 @@ async def precompute_chunk(json_input, chunk_size, current_chunk):
             if logprob >= -100000:  # Include all entries with default 0.0
                 tokens.append(get_token(entry))
                 logprobs.append(logprob)
-                top_probs = entry.get("top_logprobs", {})
-                if top_probs is None:
-                    logger.debug("top_logprobs is None for token: %s, using empty dict", get_token(entry))
-                    top_probs = {}
+                top_probs = entry.get("top_logprobs", {}) or {}
                 finite_top_probs = []
                 for key, value in top_probs.items():
                     float_value = ensure_float(value)
@@ -92,28 +87,11 @@ async def precompute_chunk(json_input, chunk_size, current_chunk):
         if start_idx >= len(tokens):
             return None, None, None
 
-        paginated_tokens = tokens[start_idx:end_idx]
-        paginated_logprobs = logprobs[start_idx:end_idx]
-        paginated_alternatives = top_alternatives[start_idx:end_idx]
-
-        return paginated_tokens, paginated_logprobs, paginated_alternatives
+        return tokens[start_idx:end_idx], logprobs[start_idx:end_idx], top_alternatives[start_idx:end_idx]
     except Exception as e:
         logger.error("Precomputation failed for chunk %d: %s", current_chunk + 1, str(e))
         return None, None, None
 
-# Synchronous wrapper for precomputation using threading
-def precompute_next_chunk_sync(json_input, current_chunk):
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    try:
-        result = loop.run_until_complete(precompute_chunk(json_input, 100, current_chunk))
-    except Exception as e:
-        logger.error("Precomputation error: %s", str(e))
-        result = None, None, None
-    finally:
-        loop.close()
-    return result
-
 # Function to process and visualize a chunk of log probs with dynamic top_logprobs
 def visualize_logprobs(json_input, chunk=0, chunk_size=100):
     try:
@@ -260,14 +238,14 @@ def visualize_logprobs(json_input, chunk=0, chunk_size=100):
 def analyze_confidence_signature(logprobs, tokens):
     if not logprobs or not tokens:
         return "No data for confidence signature analysis.", None
-    top_probs = [lps[0][1] if lps and lps[0][1] is not None else -float('inf') for lps in logprobs]  # Handle empty or None
+    top_probs = [lps[0][1] if lps and lps[0][1] is not None else -float('inf') for lps in logprobs]  # Extract top probability, handle empty or None
     if not any(p != -float('inf') for p in top_probs):
         return "No valid log probabilities for confidence analysis.", None
     moving_avg = np.convolve(top_probs, np.ones(20) / 20, mode='valid')  # 20-token window
     drops = np.where(np.diff(moving_avg) < -0.15)[0]  # Significant drops
     if not drops.size:
         return "No significant confidence drops detected.", None
-    drop_positions = [(i, tokens[i + 19] if i + 19 < len(tokens) else "End of trace") for i in drops]
+    drop_positions = [(i, tokens[i + 19] if i + 19 < len(tokens) else "End of trace") for i in drops]  # Adjust for convolution window
     return "Significant confidence drops detected at positions:", drop_positions
 
 def detect_interpretation_pivots(logprobs, tokens):
@@ -420,7 +398,7 @@ def analyze_full_trace(json_input):
 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 or visualize tokens in chunks of 100. Fixed filter ≥ -100000, dynamic number of top_logprobs, handles missing or null fields. Next chunk is precomputed proactively.")
+        gr.Markdown("Paste your JSON log prob data below to analyze reasoning traces or visualize tokens in chunks of 100. Fixed filter ≥ -100000, dynamic number of top_logprobs, handles missing or null fields.")
 
         with gr.Tabs():
             with gr.Tab("Trace Analysis"):
@@ -475,6 +453,9 @@ try:
                     outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, total_chunks_output, chunk],
                 )
 
+                def precompute_next_chunk(json_input, current_chunk):
+                    return precompute_chunk(json_input, 100, current_chunk)
+
                 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:
@@ -500,7 +481,8 @@ try:
 
                 def trigger_precomputation(json_input, current_chunk):
                     try:
-                        threading.Thread(target=precompute_next_chunk_sync, args=(json_input, current_chunk), daemon=True).start()
+                        precomputed = precompute_next_chunk(json_input, current_chunk)
+                        precomputed_next.value = precomputed  # Update state directly
                     except Exception as e:
                         logger.error("Precomputation trigger failed: %s", str(e))
                     return gr.update(value=current_chunk)