Update app.py

app.py

@@ -13,42 +13,39 @@ def load_and_preprocess_data():
 
     # Preprocess the dataset: tokenize inputs
     def preprocess_function(examples):
-        # Dynamic max length; each example will be processed to match tokenization length
         inputs = {
             'input_ids': [],
             'attention_mask': [],
             'labels': []
         }
-        for example in examples['prompt']:  # Adjust if 'prompt' is not the key for your text
+        for example in examples['prompt']:
             # Tokenize chosen and rejected parts separately
             chosen_tokenized = tokenizer(
                 example, 
                 truncation=True, 
                 padding='max_length', 
-                max_length=4  # or adjust this based on your needs
+                max_length=4
             )
             rejected_tokenized = tokenizer(
                 example, 
                 truncation=True, 
                 padding='max_length', 
-                max_length=4  # or adjust this based on your needs
+                max_length=4
             )
             
             combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
             combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
     
-            # Append inputs
+            label_length = len(combined_input_ids)
+            # Ensure labels match the exact length of input_ids
             inputs['input_ids'].append(combined_input_ids)
             inputs['attention_mask'].append(combined_attention_mask)
-            
-            # Labels should match the length of input_ids for each example
-            label_length = len(combined_input_ids)
-            inputs['labels'].append(torch.tensor([1] * (label_length // 2) + [0] * (label_length // 2)))
+            inputs['labels'].append(torch.tensor([1] * (label_length // 2) + [0] * (label_length - label_length // 2)))
     
         # 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
+        inputs['input_ids'] = torch.tensor(inputs['input_ids'], dtype=torch.long)
+        inputs['attention_mask'] = torch.tensor(inputs['attention_mask'], dtype=torch.long)
+        inputs['labels'] = torch.stack(inputs['labels'])
     
         # Debugging
         print(f"Input_ids shape: {inputs['input_ids'].shape}, Labels shape: {inputs['labels'].shape}")