Update app.py

app.py

@@ -2,14 +2,10 @@ import gradio as gr
 from datasets import load_dataset
 
 # --- Configuration ---
-DATASET_NAME = "Cnam-LMSSC/vibravox-test"
-DATASET_CONFIG = "speech_clean"
-DATASET_SPLIT = "train"
+DATASET_NAME = "Cnam-LMSSC/vibravox"
+SUBSETS = ["speech_clean", "speech_noisy", "speechless_clean", "speechless_noisy"]
+SPLITS = ["train", "validation", "test"]
 TEXT_COLUMN = "raw_text"
-
-# --- THE FINAL, CORRECT COLUMN NAMES ---
-# Based on the official dataset viewer on Hugging Face and the KeyError.
-# This list is now definitive.
 AUDIO_COLUMNS = [
     "audio.headset_microphone",
     "audio.throat_microphone",
@@ -19,60 +15,127 @@ AUDIO_COLUMNS = [
     "audio.temple_vibration_pickup"
 ]
 
-# --- Load Dataset ---
-try:
-    # Load the dataset normally.
-    dataset = load_dataset(DATASET_NAME, DATASET_CONFIG, split=DATASET_SPLIT)
-except Exception as e:
-    dataset = None
-    app_error = e
+# --- Main Application Logic ---
 
-# --- App Logic ---
-def get_audio_row(index: int):
+def load_and_update_all(subset, split):
     """
-    Retrieves a row and returns the text and the RAW audio data.
+    This is the main function. It loads a new dataset based on user selection
+    and returns updates for the entire UI, including the first row of data.
     """
-    row_index = int(index)
-    sample = dataset[row_index]
+    try:
+        # Load the newly selected dataset
+        dataset = load_dataset(DATASET_NAME, name=subset, split=split)
+        
+        # Check if the text column exists in this subset
+        has_text = TEXT_COLUMN in dataset.features
+
+        # Get the first row to display immediately
+        sample = dataset[0]
+        sentence = sample[TEXT_COLUMN] if has_text else None
+        raw_audio_data = [
+            (sample[col]['sampling_rate'], sample[col]['array']) for col in AUDIO_COLUMNS
+        ]
+        
+        # Return updates for all UI components
+        return (
+            dataset,  # Update the state object
+            gr.update(maximum=len(dataset) - 1, value=0, visible=True, interactive=True),  # Update slider
+            gr.update(value=sentence, visible=has_text),  # Update and show/hide text box
+            *raw_audio_data,  # Unpack audio data for all players
+            gr.update(value="", visible=False) # Hide any previous error messages
+        )
+    except Exception as e:
+        # If loading fails, show an error and hide the data components
+        error_message = f"Failed to load {subset}/{split}. Error: {e}"
+        empty_audio = (None, None)
+        return (
+            None, # Clear the state
+            gr.update(visible=False), # Hide slider
+            gr.update(visible=False), # Hide text box
+            *[empty_audio] * len(AUDIO_COLUMNS), # Clear all audio players
+            gr.update(value=error_message, visible=True) # Show the error message
+        )
+
 
-    sentence = sample[TEXT_COLUMN]
+def get_audio_row(dataset, index):
+    """
+    This function is called ONLY when the slider changes.
+    It fetches a new row from the currently loaded dataset (held in the state).
+    """
+    if dataset is None:
+        # This case handles when the initial load failed
+        return [None] * (1 + len(AUDIO_COLUMNS))
+
+    index = int(index)
+    sample = dataset[index]
+
+    has_text = TEXT_COLUMN in dataset.features
+    sentence = sample[TEXT_COLUMN] if has_text else None
 
-    # This will now work because we are using the correct column names.
-    # We extract the raw audio (NumPy array) and sampling rate directly.
     raw_audio_data = [
         (sample[col]['sampling_rate'], sample[col]['array']) for col in AUDIO_COLUMNS
     ]
 
     return [sentence] + raw_audio_data
 
+
 # --- Build the Gradio Interface ---
 with gr.Blocks(css="footer {display: none !important}") as demo:
-    gr.Markdown("# Vibravox Multi-Audio Viewer")
-
-    if dataset is None:
-        gr.Markdown("## 💥 Application Error")
-        gr.Markdown(f"Could not load or process the dataset. Error: `{app_error}`")
-    else:
-        gr.Markdown("Select a row to listen to all corresponding audio sensor recordings.")
+    gr.Markdown("# Vibravox Multi-Sensor Explorer")
 
-        slider = gr.Slider(minimum=0, maximum=len(dataset) - 1, step=1, value=0, label="Select Data Row")
+    # This state object holds the currently loaded dataset in memory
+    # It's invisible to the user but accessible by our functions
+    loaded_dataset_state = gr.State(None)
 
-        sentence_output = gr.Textbox(label="Raw Text", interactive=False)
+    # --- INPUT CONTROLS ---
+    with gr.Row():
+        subset_dropdown = gr.Dropdown(SUBSETS, value="speech_clean", label="Select Subset")
+        split_dropdown = gr.Dropdown(SPLITS, value="train", label="Select Split")
+    
+    # --- UI COMPONENTS FOR DATA ---
+    error_box = gr.Textbox(visible=False, interactive=False, container=False)
+    sentence_output = gr.Textbox(label="Raw Text", interactive=False, container=False)
+    slider = gr.Slider(label="Select Data Row", container=False)
+    
+    with gr.Row():
+        audio1 = gr.Audio(label="Headset Microphone")
+        audio2 = gr.Audio(label="Laryngophone (Throat Mic)")
+        audio3 = gr.Audio(label="Soft In-Ear Microphone")
+    with gr.Row():
+        audio4 = gr.Audio(label="Rigid In-Ear Microphone")
+        audio5 = gr.Audio(label="Forehead Accelerometer")
+        audio6 = gr.Audio(label="Temple Vibration Pickup")
 
-        # Labels now match the correct column names
-        with gr.Row():
-            audio1 = gr.Audio(label="Headset Microphone")
-            audio2 = gr.Audio(label="Laryngophone (Throat Mic)")
-            audio3 = gr.Audio(label="Soft In-Ear Microphone")
-        with gr.Row():
-            audio4 = gr.Audio(label="Rigid In-Ear Microphone")
-            audio5 = gr.Audio(label="Forehead Accelerometer")
-            audio6 = gr.Audio(label="Temple Vibration Pickup")
+    # A list of all the output components for easier reference
+    all_outputs = [loaded_dataset_state, slider, sentence_output, audio1, audio2, audio3, audio4, audio5, audio6, error_box]
+    audio_outputs = [sentence_output, audio1, audio2, audio3, audio4, audio5, audio6]
 
-        outputs = [sentence_output, audio1, audio2, audio3, audio4, audio5, audio6]
+    # --- WIRING THE EVENT HANDLERS ---
+    
+    # 1. When the app first loads, run the main function with default values
+    demo.load(
+        fn=load_and_update_all,
+        inputs=[subset_dropdown, split_dropdown],
+        outputs=all_outputs
+    )
+    
+    # 2. When a dropdown value changes, re-run the main function
+    subset_dropdown.change(
+        fn=load_and_update_all,
+        inputs=[subset_dropdown, split_dropdown],
+        outputs=all_outputs
+    )
+    split_dropdown.change(
+        fn=load_and_update_all,
+        inputs=[subset_dropdown, split_dropdown],
+        outputs=all_outputs
+    )
 
-        demo.load(fn=get_audio_row, inputs=gr.State(0), outputs=outputs)
-        slider.change(fn=get_audio_row, inputs=slider, outputs=outputs)
+    # 3. When ONLY the slider changes, run the simpler function
+    slider.change(
+        fn=get_audio_row,
+        inputs=[loaded_dataset_state, slider],
+        outputs=audio_outputs
+    )
 
-# Launch the application
 demo.launch()