Update app.py

app.py

@@ -232,13 +232,25 @@ def visualize_logprobs(json_input, prob_filter=-1e9, page_size=50, page=0):
         # Top 3 Token Log Probabilities (paginated)
         alt_viz_html = ""
         if paginated_logprobs and paginated_alternatives:
-            alt_viz_html = "<h3>Top 3 Token Log Probabilities (Paginated)</h3><ul>"
+            alt_viz_fig = go.Figure()
             for i, (token, probs) in enumerate(zip(paginated_tokens, paginated_alternatives)):
-                alt_viz_html += f"<li>Position {i+start_idx} (Token: {token}):<br>"
-                for tok, prob in probs:
-                    alt_viz_html += f"{tok}: {prob:.4f}<br>"
-                alt_viz_html += "</li>"
-            alt_viz_html += "</ul>"
+                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.update_layout(
+                title="Top 3 Token Log Probabilities (Paginated)",
+                xaxis_title="Token (Position)",
+                yaxis_title="Log Probability",
+                barmode='stack',
+                hovermode="closest",
+                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],
+                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, anomaly_fig, total_pages, page)
 
@@ -281,15 +293,15 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
     btn.click(
         fn=visualize_logprobs,
         inputs=[json_input, prob_filter, page_size, page],
-        outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, anomaly_output, gr.State(visible=False), gr.State(visible=False)],
+        outputs=[plot_output, table_output, text_output, alt_viz_output, drops_output, anomaly_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, visible=False)
-        current_page_output = gr.Number(label="Current Page", interactive=False, visible=False)
+        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, prob_filter, page_size, current_page, action):
         if action == "prev" and current_page > 0:
@@ -301,13 +313,13 @@ with gr.Blocks(title="Log Probability Visualizer") as app:
         return gr.update(value=current_page), gr.update(value=total_pages)
 
     prev_btn.click(
-        fn=lambda *args: update_page(*args, "prev"),
+        fn=update_page,
         inputs=[json_input, prob_filter, page_size, page, gr.State()],
         outputs=[page, total_pages_output]
     )
 
     next_btn.click(
-        fn=lambda *args: update_page(*args, "next"),
+        fn=update_page,
         inputs=[json_input, prob_filter, page_size, page, gr.State()],
         outputs=[page, total_pages_output]
     )