Update app.py

app.py

@@ -12,45 +12,45 @@ def load_and_preprocess_data():
     tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
 
     # Preprocess the dataset: tokenize inputs
-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, 
-            truncation=True, 
-            padding='max_length', 
-            max_length=max_length
-        )
-        rejected_tokenized = tokenizer(
-            example,  # Here you might need to adjust for how you handle 'chosen' and 'rejected'
-            truncation=True, 
-            padding='max_length', 
-            max_length=max_length
-        )
-        
-        combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
-        combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
-
-        # 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]
-        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'].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
+    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, 
+                truncation=True, 
+                padding='max_length', 
+                max_length=max_length
+            )
+            rejected_tokenized = tokenizer(
+                example,  # Here you might need to adjust for how you handle 'chosen' and 'rejected'
+                truncation=True, 
+                padding='max_length', 
+                max_length=max_length
+            )
+            
+            combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
+            combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
+    
+            # 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]
+            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'].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)