Update app.py

app.py

@@ -12,27 +12,32 @@ def load_and_preprocess_data():
     tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
 
     # 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
+def preprocess_function(examples):
+    # Assuming 'examples' is a batch of examples
+    max_length = 500  # Half of 1000 for each segment
+    inputs = {
+        'input_ids': [],
+        'attention_mask': [],
+        'labels': []
+    }
+    for example in examples['prompt']:  # Adjust if 'prompt' is not the key for your text
         chosen_tokenized = tokenizer(
-            example['prompt'], 
+            example, 
             truncation=True, 
             padding='max_length', 
             max_length=max_length
         )
         rejected_tokenized = tokenizer(
-            example['prompt'], 
+            example,  # Here you might need to adjust for how you handle 'chosen' and 'rejected'
             truncation=True, 
             padding='max_length', 
             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
+
+        # Pad or truncate to ensure length is 1000
         if len(combined_input_ids) > 1000:
             combined_input_ids = combined_input_ids[:1000]
             combined_attention_mask = combined_attention_mask[:1000]
@@ -40,17 +45,12 @@ def load_and_preprocess_data():
             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': combined_input_ids,
-            'attention_mask': combined_attention_mask,
-            'labels': [1] * len(example['chosen']) + [0] * len(example['rejected'])
-        }
-        
-        # Truncate labels if they exceed the length of input_ids
-        inputs['labels'] = inputs['labels'][:len(inputs['input_ids'])]
-        
-        return inputs
+
+        inputs['input_ids'].append(combined_input_ids)
+        inputs['attention_mask'].append(combined_attention_mask)
+        inputs['labels'].append([1] * 500 + [0] * 500)  # Adjust based on your label structure
+
+    return inputs
     # Apply preprocessing to the dataset
     dataset = dataset.map(preprocess_function, batched=True)