Upload app.py

app.py

@@ -858,21 +858,30 @@ def create_statistics_dashboard(df: pd.DataFrame) -> Tuple[str, go.Figure]:
 current_data = {"df": pd.DataFrame(), "type": "unknown"}
 
 
-def load_dataset_action(source_type: str, dataset_id: str, file_upload) -> Tuple[str, str]:
-    """Handle dataset loading."""
+def load_dataset_action(source_type: str, dataset_id: str, file_upload):
+    """Handle dataset loading and return all visualization updates."""
     global current_data
 
     if source_type == "HuggingFace Hub":
         if not dataset_id:
-            return "Please enter a dataset ID", ""
+            empty_fig = go.Figure()
+            empty_fig.update_layout(template="plotly_dark")
+            return ("Please enter a dataset ID", "", "No data", empty_fig, empty_fig, empty_fig, "No data", empty_fig,
+                    gr.update(maximum=0), gr.update(choices=[], value=None))
         df, msg = load_hf_dataset(dataset_id)
     else:  # Local File
         if file_upload is None:
-            return "Please upload a file", ""
+            empty_fig = go.Figure()
+            empty_fig.update_layout(template="plotly_dark")
+            return ("Please upload a file", "", "No data", empty_fig, empty_fig, empty_fig, "No data", empty_fig,
+                    gr.update(maximum=0), gr.update(choices=[], value=None))
         df, msg = load_jsonl_file(file_upload.name)
 
     if df.empty:
-        return msg, ""
+        empty_fig = go.Figure()
+        empty_fig.update_layout(template="plotly_dark")
+        return (msg, "", "No data", empty_fig, empty_fig, empty_fig, "No data", empty_fig,
+                gr.update(maximum=0), gr.update(choices=[], value=None))
 
     current_data["df"] = df
     current_data["type"] = detect_dataset_type(df)
@@ -881,7 +890,27 @@ def load_dataset_action(source_type: str, dataset_id: str, file_upload) -> Tuple
     if len(df.columns) > 10:
         columns_info += f" ... and {len(df.columns) - 10} more"
 
-    return msg, f"Dataset type: {current_data['type']}\n{columns_info}"
+    # Generate all visualizations
+    stats_html, stats_fig = create_statistics_dashboard(df)
+    graph_fig = create_thought_anchor_graph(df)
+    embed_fig = create_embedding_visualization(df)
+    circuit_html, circuit_fig = create_circuit_visualization(df)
+
+    # Generate query list
+    query_choices = []
+    if 'query' in df.columns:
+        queries = df['query'].unique().tolist()
+        for i, q in enumerate(queries):
+            q_str = str(q) if q is not None else ""
+            if len(q_str) > 80:
+                query_choices.append(f"[{i+1}] {q_str[:77]}...")
+            else:
+                query_choices.append(f"[{i+1}] {q_str}")
+
+    return (msg, f"Dataset type: {current_data['type']}\n{columns_info}",
+            stats_html, stats_fig, graph_fig, embed_fig, circuit_html, circuit_fig,
+            gr.update(maximum=max(0, len(df) - 1)),
+            gr.update(choices=query_choices, value=None))
 
 
 def get_token_details(idx: int) -> Tuple[str, go.Figure]:
@@ -1180,19 +1209,7 @@ with gr.Blocks(title="PTS Visualizer", css=CSS) as demo:
     load_btn.click(
         fn=load_dataset_action,
         inputs=[source_type, dataset_dropdown, file_upload],
-        outputs=[load_status, dataset_info],
-        api_name=False
-    ).then(
-        fn=refresh_all,
-        outputs=[stats_html, stats_chart, graph_plot, embed_plot, circuit_html, circuit_chart],
-        api_name=False
-    ).then(
-        fn=lambda: gr.update(maximum=max(0, len(current_data["df"]) - 1)),
-        outputs=[token_slider],
-        api_name=False
-    ).then(
-        fn=get_query_list,
-        outputs=[query_filter],
+        outputs=[load_status, dataset_info, stats_html, stats_chart, graph_plot, embed_plot, circuit_html, circuit_chart, token_slider, query_filter],
         api_name=False
     )