Update app.py

app.py

@@ -2,13 +2,13 @@ import gradio as gr
 from datasets import load_dataset
 
 # --- Configuration ---
-# The ONLY change is on this line: we're pointing back to the smaller test dataset.
 DATASET_NAME = "Cnam-LMSSC/vibravox-test"
-# ---------------------------------------------------------------------------------
-
 SUBSETS = ["speech_clean", "speech_noisy", "speechless_clean", "speechless_noisy"]
 SPLITS = ["train", "validation", "test"]
 TEXT_COLUMN = "raw_text"
+# Add new column names to the configuration
+PHONEMIZED_TEXT_COLUMN = "phonemized_text"
+GENDER_COLUMN = "gender"
 AUDIO_COLUMNS = [
     "audio.headset_microphone",
     "audio.throat_microphone",
@@ -22,82 +22,88 @@ AUDIO_COLUMNS = [
 
 def load_and_update_all(subset, split):
     """
-    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.
+    Loads a new dataset and returns updates for the entire UI.
     """
     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
+        has_text_fields = TEXT_COLUMN in dataset.features
 
         # Get the first row to display immediately
         sample = dataset[0]
-        sentence = sample[TEXT_COLUMN] if has_text else None
+        sentence = sample[TEXT_COLUMN] if has_text_fields else None
+        # Fetch the new fields
+        phonemized_text = sample[PHONEMIZED_TEXT_COLUMN] if has_text_fields else None
+        gender = sample[GENDER_COLUMN] if has_text_fields else None
+        
         raw_audio_data = [
             (sample[col]['sampling_rate'], sample[col]['array']) for col in AUDIO_COLUMNS
         ]
         
-        # Return updates for all UI components
+        # Return updates for all UI components, including the new ones
         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
+            dataset,
+            gr.update(maximum=len(dataset) - 1, value=0, visible=True, interactive=True),
+            gr.update(value=sentence, visible=has_text_fields),
+            # Add updates for the new text boxes
+            gr.update(value=phonemized_text, visible=has_text_fields),
+            gr.update(value=gender, visible=has_text_fields),
+            *raw_audio_data,
+            gr.update(value="", visible=False)
         )
     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 empty/hidden updates for all components on error
         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
+            None,
+            gr.update(visible=False),
+            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
+            *[empty_audio] * len(AUDIO_COLUMNS),
+            gr.update(value=error_message, visible=True)
         )
 
-
 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).
+    Fetches a new row from the currently loaded dataset when the slider moves.
     """
     if dataset is None:
-        # This case handles when the initial load failed
-        return [None] * (1 + len(AUDIO_COLUMNS))
+        return [None] * (3 + len(AUDIO_COLUMNS)) # 3 text fields now
 
     index = int(index)
     sample = dataset[index]
 
-    has_text = TEXT_COLUMN in dataset.features
-    sentence = sample[TEXT_COLUMN] if has_text else None
+    has_text_fields = TEXT_COLUMN in dataset.features
+    sentence = sample[TEXT_COLUMN] if has_text_fields else None
+    # Fetch the new fields for the selected row
+    phonemized_text = sample[PHONEMIZED_TEXT_COLUMN] if has_text_fields else None
+    gender = sample[GENDER_COLUMN] if has_text_fields else None
 
     raw_audio_data = [
         (sample[col]['sampling_rate'], sample[col]['array']) for col in AUDIO_COLUMNS
     ]
 
-    return [sentence] + raw_audio_data
-
+    return [sentence, phonemized_text, gender] + raw_audio_data
 
 # --- Build the Gradio Interface ---
 with gr.Blocks(css="footer {display: none !important}") as demo:
-    gr.Markdown("# Vibravox Multi-Sensor Explorer (Test Dataset)")
+    # Change the app title
+    gr.Markdown("# Vibravox Viewer")
 
-    # This state object holds the currently loaded dataset in memory
     loaded_dataset_state = gr.State(None)
 
-    # --- 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)
+    
+    # Group the text outputs together
+    with gr.Row():
+        sentence_output = gr.Textbox(label="Raw Text", interactive=False)
+        phonemized_output = gr.Textbox(label="Phonemized Text", interactive=False)
+        gender_output = gr.Textbox(label="Gender", interactive=False)
+
+    slider = gr.Slider(label="Select Data Row")
     
     with gr.Row():
         audio1 = gr.Audio(label="Headset Microphone")
@@ -108,29 +114,21 @@ with gr.Blocks(css="footer {display: none !important}") as demo:
         audio5 = gr.Audio(label="Forehead Accelerometer")
         audio6 = gr.Audio(label="Temple Vibration Pickup")
 
-    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]
+    # Update the component lists to include the new text boxes
+    all_outputs = [loaded_dataset_state, slider, sentence_output, phonemized_output, gender_output, audio1, audio2, audio3, audio4, audio5, audio6, error_box]
+    data_outputs = [sentence_output, phonemized_output, gender_output, audio1, audio2, audio3, audio4, audio5, audio6]
 
     # --- WIRING THE EVENT HANDLERS ---
-    demo.load(
-        fn=load_and_update_all,
-        inputs=[subset_dropdown, split_dropdown],
-        outputs=all_outputs
-    )
-    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
-    )
-    slider.change(
-        fn=get_audio_row,
-        inputs=[loaded_dataset_state, slider],
-        outputs=audio_outputs
-    )
+    # The handlers themselves don't need to change, as we updated the functions and component lists
+    
+    # 1. When the app first loads
+    demo.load(fn=load_and_update_all, inputs=[subset_dropdown, split_dropdown], outputs=all_outputs)
+    
+    # 2. When a dropdown value changes
+    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)
+
+    # 3. When ONLY the slider changes
+    slider.change(fn=get_audio_row, inputs=[loaded_dataset_state, slider], outputs=data_outputs)
 
 demo.launch()