Update app.py

app.py

@@ -13,7 +13,7 @@ def load_and_preprocess_data():
 
     # Preprocess the dataset: tokenize inputs
     def preprocess_function(examples):
-        max_length = 4  # Now using 4 for both chosen and rejected
+        max_length = 4  # Using 4 for both chosen and rejected
         inputs = {
             'input_ids': [],
             'attention_mask': [],
@@ -36,17 +36,18 @@ def load_and_preprocess_data():
             combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
             combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
     
-            # No need for padding or truncating since we're already at max length
             inputs['input_ids'].append(combined_input_ids)
             inputs['attention_mask'].append(combined_attention_mask)
-            # Adjust labels to match new length; here assuming binary classification where first half is 'chosen'
-            inputs['labels'].append(torch.tensor([1] * 2 + [0] * 2))  # 4 labels for 4 tokens
+            # Labels should be exactly 8 for each example (4 for chosen, 4 for rejected)
+            inputs['labels'].append(torch.tensor([1] * 4 + [0] * 4))  # 8 labels for 8 tokens
     
         # Convert lists to tensors for each batch
         inputs['input_ids'] = torch.tensor(inputs['input_ids'])
         inputs['attention_mask'] = torch.tensor(inputs['attention_mask'])
         inputs['labels'] = torch.stack(inputs['labels'])  # Stack labels along batch dimension
     
+        print(f"Input_ids shape: {inputs['input_ids'].shape}, Labels shape: {inputs['labels'].shape}")  # Debugging
+    
         return inputs
     # Apply preprocessing to the dataset
     dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)