Add HF Dataset subset and split support

app.py

@@ -10,10 +10,11 @@ from src.visualization import generate_charts
 from src.llm import get_insights, get_followup_questions, ask_llm, get_ml_recommendations, analyze_text_content
 
 # Updated analyze_dataset to accept api_token
-def analyze_dataset(file_obj, hf_dataset_name, api_token):
+# Updated analyze_dataset to accept subset and split
+def analyze_dataset(file_obj, hf_dataset_name, hf_subset, hf_split, api_token):
     # Handle HF Dataset loading if name is provided
     if hf_dataset_name:
-        df, error, load_log = load_hf_dataset(hf_dataset_name, api_token=api_token)
+        df, error, load_log = load_hf_dataset(hf_dataset_name, subset=hf_subset, split=hf_split, api_token=api_token)
     elif file_obj:
         df, error, load_log = load_file(file_obj)
     else:
@@ -104,8 +105,9 @@ def chat_response(message, history, overview_text, api_token):
     return ask_llm(message, history, overview_text, api_token)
 
 # Updated process function wrapper to match inputs/outputs
-def process_data_wrapper(file_obj, hf_dataset, api_token):
-    results = analyze_dataset(file_obj, hf_dataset, api_token)
+# Updated process function wrapper to match inputs/outputs
+def process_data_wrapper(file_obj, hf_dataset, hf_subset, hf_split, api_token):
+    results = analyze_dataset(file_obj, hf_dataset, hf_subset, hf_split, api_token)
     return results + (results[0],) # Append overview_md for the state
 
 # --- Custom Styling & Theme ---
@@ -250,7 +252,10 @@ with gr.Blocks(title="Auto Data Analyst", theme=theme, css=custom_css) as demo:
                     example_btn = gr.Button("🎲 Load Sample Data", variant="secondary")
 
                 with gr.TabItem("HF Dataset"):
-                    hf_input = gr.Textbox(label="Dataset Name", placeholder="e.g. titanic, dair-ai/emotion", info="Loads the 'train' split.")
+                    hf_input = gr.Textbox(label="Dataset Name", placeholder="e.g. glue", info="Name of the dataset on Hub.")
+                    with gr.Row():
+                        hf_subset_input = gr.Textbox(label="Subset/Config (Optional)", placeholder="e.g. mrpc", info="Specific configuration.", scale=1)
+                        hf_split_input = gr.Textbox(label="Split", value="train", placeholder="e.g. train, test", info="Split to load.", scale=1)
                     hf_load_btn = gr.Button("⬇️ Load Dataset", variant="primary")
 
             gr.Markdown("---")
@@ -297,7 +302,7 @@ with gr.Blocks(title="Auto Data Analyst", theme=theme, css=custom_css) as demo:
     # Event wiring - File Upload
     file_upload.change(
         fn=process_data_wrapper,
-        inputs=[file_upload, gr.State(None), api_token_input],
+        inputs=[file_upload, gr.State(None), gr.State(None), gr.State(None), api_token_input],
         outputs=[
             overview_md,
             dataframe_view,
@@ -315,7 +320,7 @@ with gr.Blocks(title="Auto Data Analyst", theme=theme, css=custom_css) as demo:
     # Also wire HF Load Button
     hf_load_btn.click(
         fn=process_data_wrapper,
-        inputs=[gr.State(None), hf_input, api_token_input],
+        inputs=[gr.State(None), hf_input, hf_subset_input, hf_split_input, api_token_input],
         outputs=[
             overview_md,
             dataframe_view,

src/ingestion.py

@@ -123,7 +123,7 @@ def load_file(file_obj):
     except Exception as e:
         return None, f"Error loading file: {str(e)}", None
 
-def load_hf_dataset(dataset_name, split='train', api_token=None):
+def load_hf_dataset(dataset_name, subset=None, split='train', api_token=None):
     """
     Loads a dataset from Hugging Face Hub.
     """
@@ -131,17 +131,19 @@ def load_hf_dataset(dataset_name, split='train', api_token=None):
         from datasets import load_dataset
         
         # Load dataset
-        # If config is needed, user might need to specify "dataset_name/config". 
-        # For now, we try default.
-        ds = load_dataset(dataset_name, split=split, token=api_token, trust_remote_code=True)
+        # If subset is provided, pass it as the second argument
+        if subset:
+            ds = load_dataset(dataset_name, subset, split=split, token=api_token, trust_remote_code=True)
+        else:
+            ds = load_dataset(dataset_name, split=split, token=api_token, trust_remote_code=True)
         
         # Convert to pandas
         df = ds.to_pandas()
         
         if df.empty:
-             return None, f"Dataset '{dataset_name}' (split='{split}') is empty.", None
+             return None, f"Dataset '{dataset_name}' (subset='{subset}', split='{split}') is empty.", None
             
-        return df, None, f"Loaded Hugging Face Dataset: `{dataset_name}` (Split: {split})"
+        return df, None, f"Loaded Hugging Face Dataset: `{dataset_name}`\n- Subset: `{subset}`\n- Split: `{split}`"
         
     except Exception as e:
         return None, f"Error loading HF Dataset '{dataset_name}': {str(e)}", None

verify_pipeline_mock.py

@@ -27,8 +27,8 @@ with patch('src.llm.get_insights', return_value="Mocked Insights") as mock_insig
     print(f"Running pipeline with MOCKED LLM...{os.path.basename(mock_file.name)}")
     
     try:
-        # analyze_dataset signature changed: (file_obj, hf_dataset_name, api_token)
-        results = analyze_dataset(mock_file, None, api_token="test")
+        # analyze_dataset signature changed: (file_obj, hf_dataset_name, hf_subset, hf_split, api_token)
+        results = analyze_dataset(mock_file, None, None, None, api_token="test")
         
         # Unpack results to verify types (updated for new return signature)
         # (overview_output, df_head, insights, chart, anomaly_md, anomalies_df, ml_recs, text_analysis_output, download_path)