Update app.py

app.py

@@ -10,9 +10,11 @@ def load_and_preprocess_data():
 
     # Extract "prompt" and "chosen"/"rejected" as input-output pairs
     def process_row(example):
+         row = example['row']  # Access the 'row' key inside the example
         return {
-            'prompt': example['row']['prompt'],
-            'response': example['row']['chosen'],  # You can change this to 'rejected' if needed
+            'prompt': row['prompt'],
+            'chosen': row['chosen'],
+            'rejected': row['rejected'],
         }
     
     dataset = dataset.map(process_row)