Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from datasets import load_dataset, Features, Value, Audio, Dataset
 from huggingface_hub import HfApi, create_repo
-import pandas as pd  # Import pandas for displaying the dataset
+import pandas as pd
 
 
 def filter_dataset(dataset_name, split_name, keywords_text):
@@ -24,12 +24,11 @@ def filter_dataset(dataset_name, split_name, keywords_text):
         data_for_df = []  # Store data for DataFrame
         for i, example in enumerate(filtered_dataset):
             matching_indices.append(i)
-            #  Extract data and append.  Crucially, *decode* audio here.
+            #  Extract data and append. Handle potential KeyErrors.
             example_data = {
-                'prompt': example['prompt'],
-                'strategy': example['strategy'],
-                'seed': example['seed'],
-                'audio': example['audio']['array']  # Get the NumPy array
+                'prompt': example.get('prompt', None),  # Use .get() for safety
+                'chosen': example.get('chosen', {}).get('array', None) if isinstance(example.get('chosen'), dict) else None, # Handle nested structure, check if it's a dict
+                'rejected': example.get('rejected', {}).get('array', None) if isinstance(example.get('rejected'), dict) else None,  # Handle nested structure
             }
             data_for_df.append(example_data)
 
@@ -47,37 +46,32 @@ def filter_dataset(dataset_name, split_name, keywords_text):
 def push_to_hub(df_json, dataset_name, split_name, new_dataset_repo_id, hf_token):
     """Pushes a Pandas DataFrame (from JSON) to the Hugging Face Hub."""
     if not hf_token:
-        return "Error: Hugging Face token is required.", None
+        return "Error: Hugging Face Token is required.", None
 
     try:
         # Convert JSON back to DataFrame
         df = pd.read_json(df_json)
-        
+
         if df.empty:
-            return "Error: Cannot push an empty dataset",None
-            
+            return "Error: Cannot push an empty dataset", None
+
         # Convert DataFrame to Hugging Face Dataset
         dataset = Dataset.from_pandas(df)
 
-        # --- Load original (for feature definition)
-        full_dataset = load_dataset(dataset_name, split=split_name, streaming=False)
-        
-        if len(full_dataset) == 0:
-            return "Error: Source Dataset Appears Empty",None
 
         # --- 5. Define features (for consistent schema) ---
-        features = Features({
+        features_dict = {
             'prompt': Value(dtype='string', id=None),
-            'audio': Audio(sampling_rate=16000),
-            'strategy': Value(dtype='string', id=None),
-            'seed': Value(dtype='int64', id=None)
-        })
-        
+            'chosen': Audio(sampling_rate=16000), # Assuming 16kHz; adjust if needed
+            'rejected': Audio(sampling_rate=16000), # Assuming 16kHz
+        }
+
+        features = Features(features_dict)
+
         try:
-          dataset = dataset.cast(features)
+            dataset = dataset.cast(features)
         except Exception as e:
-            return f"An error occurred: {e}",None
-
+            return f"An error occurred during casting: {e}", None
 
         # --- 6. Upload to the Hugging Face Hub ---
         api = HfApi(token=hf_token)
@@ -94,14 +88,12 @@ def push_to_hub(df_json, dataset_name, split_name, new_dataset_repo_id, hf_token
 
     except Exception as e:
         return f"An error occurred during push: {e}", None
-
-
 # --- Gradio Interface ---
 with gr.Blocks() as demo:
     gr.Markdown("# Dataset Filter and Push")
 
     with gr.Row():
-        dataset_name_input = gr.Textbox(label="Source Dataset Name", value="declare-lab/audio-alpaca")
+        dataset_name_input = gr.Textbox(label="Source Dataset Name", value="ashraq/esc50") # Example with chosen/rejected
         split_name_input = gr.Textbox(label="Split Name", value="train")
 
     keywords_input = gr.Textbox(label="Keywords (comma-separated)", value="dog, cat")