Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import pandas as pd
-from datasets import load_dataset, get_dataset_split_names
+from datasets import load_dataset, get_dataset_split_names, get_dataset_config_info
 from huggingface_hub import HfApi
 import os
 import pathlib
@@ -45,28 +45,33 @@ def get_dataset_splits(dataset_id: str):
     if not dataset_id:
         return gr.update(choices=[], value=None, interactive=False)
     try:
-        splits = get_dataset_split_names(dataset_id)
-        # Set the first split as the default value to trigger the next event
+        splits = get_dataset_split_names(dataset_id, trust_remote_code=True)
         return gr.update(choices=splits, value=splits[0] if splits else None, interactive=True)
     except Exception as e:
         gr.Warning(f"Could not fetch splits for dataset '{dataset_id}'. Error: {e}")
         return gr.update(choices=[], value=None, interactive=False)
 
-def get_split_columns(dataset_id: str, split: str):
-    """Gets all columns for a selected split by loading its metadata."""
-    if not dataset_id or not split:
+def get_split_columns(dataset_id: str):
+    """Gets all columns for a selected dataset by loading its metadata info."""
+    if not dataset_id:
         return gr.update(choices=[], value=None, interactive=False)
     try:
-        # Get the .features property from the dataset info.
-        features = load_dataset(dataset_id, split=split, streaming=True).features
+        # --- THIS IS THE ROBUST FIX ---
+        # Use get_dataset_config_info to get schema without loading data.
+        # This is the official and most reliable way.
+        info = get_dataset_config_info(dataset_id, trust_remote_code=True)
+        features = info.features
+        
+        # The user is right, we should show ALL columns.
         columns = list(features.keys())
         
-        # Heuristically find the best text column
+        # We can still be helpful by guessing the best default.
         preferred_cols = ['text', 'content', 'instruction', 'question', 'document', 'prompt']
         best_col = next((col for col in preferred_cols if col in columns), columns[0] if columns else None)
+        
         return gr.update(choices=columns, value=best_col, interactive=True)
     except Exception as e:
-        gr.Warning(f"Could not fetch columns for split '{split}'. Error: {e}")
+        gr.Warning(f"Could not get columns for '{dataset_id}'. It might be a gated dataset or have an unusual structure. Error: {e}")
         return gr.update(choices=[], value=None, interactive=False)
 
 # --- Main Atlas Generation Logic ---
@@ -88,7 +93,7 @@ def generate_atlas(
         
     progress(0, desc=f"Loading dataset '{dataset_name}' [{split}]...")
     try:
-        dataset = load_dataset(dataset_name, split=split)
+        dataset = load_dataset(dataset_name, split=split, trust_remote_code=True)
         df = dataset.to_pandas()
     except Exception as e:
         raise gr.Error(f"Failed to load data. Error: {e}")
@@ -163,14 +168,14 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Embedding Atlas Explorer") as app:
             output_html = gr.HTML("<div style='display:flex; justify-content:center; align-items:center; height:800px; border: 1px solid #ddd; border-radius: 5px;'><p>Atlas will be displayed here after generation.</p></div>")
 
     # --- Chained Event Listeners for Dynamic UI ---
-    # When the user submits a name, get their datasets
     hf_user_input.submit(fn=get_user_datasets, inputs=hf_user_input, outputs=dataset_input)
     
-    # --- THIS IS THE FIX ---
-    # Use .change() so that when a dataset is selected (by user OR another function), it triggers the next step.
+    # When a dataset is selected, get its splits.
     dataset_input.change(fn=get_dataset_splits, inputs=dataset_input, outputs=split_input)
-    split_input.change(fn=get_split_columns, inputs=[dataset_input, split_input], outputs=text_column_input)
-
+    
+    # When a dataset is selected, ALSO get its columns. The split doesn't matter for column schema.
+    dataset_input.change(fn=get_split_columns, inputs=dataset_input, outputs=text_column_input)
+    
     # --- Button Click Event ---
     generate_button.click(
         fn=generate_atlas,