Update app.py

app.py

@@ -12,41 +12,46 @@ def load_and_preprocess_data():
     tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
 
     # Preprocess the dataset: tokenize inputs
-    def preprocess_function(examples):
-        max_length = 4  # Using 4 for both chosen and rejected
+   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
+            # Tokenize chosen and rejected parts separately
             chosen_tokenized = tokenizer(
                 example, 
                 truncation=True, 
                 padding='max_length', 
-                max_length=max_length
+                max_length=4  # or adjust this based on your needs
             )
             rejected_tokenized = tokenizer(
-                example,  
+                example, 
                 truncation=True, 
                 padding='max_length', 
-                max_length=max_length
+                max_length=4  # or adjust this based on your needs
             )
             
             combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
             combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
     
+            # Append inputs
             inputs['input_ids'].append(combined_input_ids)
             inputs['attention_mask'].append(combined_attention_mask)
-            # 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
+            
+            # 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)))
     
         # 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
+        # Debugging
+        print(f"Input_ids shape: {inputs['input_ids'].shape}, Labels shape: {inputs['labels'].shape}")
     
         return inputs
     # Apply preprocessing to the dataset