Update app.py

app.py

@@ -10,14 +10,31 @@ def load_and_preprocess_data():
 
     # Extract "prompt" and "chosen"/"rejected" as input-output pairs
     def process_row(example):
-        row = example['rows']  # Access the 'row' key inside the example
+        # Access the 'rows' key
+        rows = example['rows']
+        
+        # Initialize lists for prompts, chosen, and rejected pairs
+        prompts = []
+        chosen_responses = []
+        rejected_responses = []
+        
+        # Process each row within 'rows'
+        for row_data in rows:
+            row = row_data['row']  # Access the 'row' inside each row_data
+        
+            prompts.append(row['prompt'])
+            chosen_responses.append(row['chosen'])
+            rejected_responses.append(row['rejected'])
+        
+        # Return the processed result in the correct format
         return {
-            'prompt': row['prompt'],
-            'chosen': row['chosen'],
-            'rejected': row['rejected'],
+            'prompt': prompts,
+            'chosen': chosen_responses,
+            'rejected': rejected_responses
         }
-    
-    dataset = dataset.map(process_row)
+
+    # Apply the processing function to the dataset
+    dataset = dataset.map(process_row, batched=True)
 
     # Load the tokenizer (You can choose any model, here we use BERT)
     tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')