Update app.py

app.py

@@ -14,35 +14,43 @@ def load_and_preprocess_data():
     # Preprocess the dataset: tokenize inputs
     def preprocess_function(example):
         # Tokenize both chosen and rejected responses with a larger max_length
+        max_length = 500  # Adjust so combined doesn't exceed 1000
         chosen_tokenized = tokenizer(
             example['prompt'], 
             truncation=True, 
             padding='max_length', 
-            max_length=2000  # Adjust based on your longest example but within model limits
+            max_length=max_length
         )
         rejected_tokenized = tokenizer(
             example['prompt'], 
             truncation=True, 
             padding='max_length', 
-            max_length=2000  # Adjust based on your longest example but within model limits
+            max_length=max_length
         )
         
         # Combine chosen and rejected in one sequence
+        combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
+        combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
+    
+        # Ensure the combined length does not exceed 1000 tokens
+        if len(combined_input_ids) > 1000:
+            combined_input_ids = combined_input_ids[:1000]
+            combined_attention_mask = combined_attention_mask[:1000]
+        else:
+            padding_length = 1000 - len(combined_input_ids)
+            combined_input_ids += [tokenizer.pad_token_id] * padding_length
+            combined_attention_mask += [0] * padding_length
+    
         inputs = {
-            'input_ids': chosen_tokenized['input_ids'] + rejected_tokenized['input_ids'],
-            'attention_mask': chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask'],
+            'input_ids': combined_input_ids,
+            'attention_mask': combined_attention_mask,
             'labels': [1] * len(example['chosen']) + [0] * len(example['rejected'])
         }
         
-        # Ensure the lengths match; adjust if necessary
-        max_length = 4096  # Longformer's max length
-        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]
-
+        # Truncate labels if they exceed the length of input_ids
+        inputs['labels'] = inputs['labels'][:len(inputs['input_ids'])]
+        
         return inputs
-
     # Apply preprocessing to the dataset
     dataset = dataset.map(preprocess_function, batched=True)