Update app.py

app.py

@@ -1,29 +1,52 @@
 # Import necessary libraries
 import gradio as gr
+from typing import Dict, Any
 from data_loader import read_dask_data, read_polars_data, read_another_dask_data
 
-# Function to handle the data loading logic
-def load_and_display_data():
-    # Load datasets
-    dask_data = read_dask_data()
-    polars_data = read_polars_data()
-    another_dask_data = read_another_dask_data()
+# Function to handle the data loading and processing logic
+def load_and_process_data(dataset_choice: str, num_rows: int) -> Dict[str, Any]:
+    try:
+        # Load and process datasets based on user choice
+        if dataset_choice == "Dask Data":
+            data = read_dask_data()
+            processed_data = data.head(num_rows)
+        elif dataset_choice == "Polars Data":
+            data = read_polars_data()
+            processed_data = data.head(num_rows)
+        elif dataset_choice == "Another Dask Data":
+            data = read_another_dask_data()
+            processed_data = data.head(num_rows)
+        else:
+            return {"error": "Invalid dataset choice."}
 
-    # Optionally process data (e.g., filter, summarize)
-    # Display the first few rows of the datasets
-    return {
-        "Dask Data from Codeforces Python Submissions": dask_data.head(),
-        "Polars Data from Python Reasoning Dataset": polars_data.head(),
-        "Another Dask Data from Python Github Code": another_dask_data.head()
-    }
+        # Optionally, include more complex data processing here
 
-# Create a Gradio interface to display the data
-demo = gr.Interface(
-    load_and_display_data,
-    inputs=[],  # No inputs for this simple demo
-    outputs="json",  # Display output as a JSON to show data head
-    title="Dataset Loader Demo",
-)
+        # Return the processed data
+        return {
+            "Processed Data": processed_data
+        }
+    except Exception as e:
+        # Log the exception
+        print(f"Error processing data: {str(e)}")
+        return {"error": "Unable to process data. Please check the logs for details."}
+
+# Create a Gradio interface with more features
+def create_interface():
+    # Interface inputs
+    dataset_choice = gr.inputs.Dropdown(["Dask Data", "Polars Data", "Another Dask Data"], label="Select Dataset")
+    num_rows = gr.inputs.Slider(minimum=1, maximum=100, default=5, label="Number of Rows to Display")
+
+    # Interface definition
+    demo = gr.Interface(
+        fn=load_and_process_data,
+        inputs=[dataset_choice, num_rows],  # User can select dataset and number of rows
+        outputs="json",  # Display output as JSON
+        title="Enhanced Dataset Loader Demo",
+        description="Interact with various datasets and select the amount of data to display.",
+        allow_flagging="never"  # Disable flagging if not needed
+    )
 
-if __name__ == "__main__":
     demo.launch()
+
+if __name__ == "__main__":
+    create_interface()