fix update_columns

app.py

@@ -122,22 +122,26 @@ def train_model(dataset_id, custom_dataset_id, label_column, feature_columns,
 
     return output_text, fig, columns_available
 
-def update_columns(dataset_id, custom_dataset_id):
+def update_columns(dataset_id, dataset_config, custom_dataset_id):
     """
-    Callback to dynamically fetch the columns from the dataset
-    so the user can pick which columns to use as features/labels.
+    Load the dataset from HF hub, using either the suggested one or the custom user-specified,
+    plus an optional config. 
     """
     if dataset_id != "SKIP/ENTER_CUSTOM":
         final_id = dataset_id
+        final_config = dataset_config.strip() if dataset_config else None
     else:
+        # Use the user-supplied text
         final_id = custom_dataset_id.strip()
+        final_config = None  # or parse from text if you like
 
-    # Try to load the dataset and return columns
     try:
-        ds = load_dataset(final_id, split="train")
+        if final_config:
+            ds = load_dataset(final_id, final_config, split="train")
+        else:
+            ds = load_dataset(final_id, split="train")
         df = pd.DataFrame(ds)
         cols = df.columns.tolist()
-        # Return as list of selectable options
         return gr.update(choices=cols), gr.update(choices=cols), f"Columns found: {cols}"
     except Exception as e:
         return gr.update(choices=[]), gr.update(choices=[]), f"Error loading {final_id}: {e}"