Update app.py

app.py

@@ -12,23 +12,38 @@ def load_and_preprocess_data():
     # Load the tokenizer
     tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
 
-    # Preprocess the dataset: tokenize inputs
+    # Preprocess the dataset: tokenize 
     def preprocess_function(example):
-        inputs = tokenizer(
+        # Tokenize both chosen and rejected responses
+        chosen_tokenized = tokenizer(
             example['prompt'], 
             truncation=True, 
             padding='max_length', 
             max_length=512
         )
-        # Assume binary classification: 1 for chosen, 0 for rejected
-        inputs['labels'] = [1] * len(example['chosen']) + [0] * len(example['rejected'])
-        inputs['input_ids'] = (
-            inputs['input_ids'][:len(example['chosen'])] + 
-            inputs['input_ids'][len(example['chosen']):]
+        rejected_tokenized = tokenizer(
+            example['prompt'], 
+            truncation=True, 
+            padding='max_length', 
+            max_length=512
         )
+        
+        # Combine chosen and rejected in one sequence, but you'll need to ensure they fit within max_length
+        # Here's a simplified approach; adjust based on your needs
+        inputs = {
+            'input_ids': chosen_tokenized['input_ids'] + rejected_tokenized['input_ids'],
+            'attention_mask': chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask'],
+            'labels': [1] * len(example['chosen']) + [0] * len(example['rejected'])
+        }
+        
+        # Ensure the lengths match; you might need to truncate or pad
+        max_length = 512 * 2  # Since we're combining two sequences
+        if len(inputs['input_ids']) > max_length:
+            inputs['input_ids'] = inputs['input_ids'][:max_length]
+            inputs['attention_mask'] = inputs['attention_mask'][:max_length]
+            inputs['labels'] = inputs['labels'][:max_length]  # Be cautious here, might need adjustment based on actual label length
+    
         return inputs
-
-    # Apply preprocessing to the dataset
     dataset = dataset.map(preprocess_function, batched=True)
 
     # Convert dataset to PyTorch format