Update app.py

app.py

@@ -5,13 +5,10 @@ import torch
 
 # Load and preprocess dataset
 def load_and_preprocess_data():
-    # Load the dataset from the uploaded JSON file
-    dataset = load_dataset('json', data_files='data.json')['train']  # Assuming single split
+    dataset = load_dataset('json', data_files='data.json')['train']
 
-    # Load the tokenizer
     tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
 
-    # Preprocess the dataset: tokenize inputs
     def preprocess_function(examples):
         inputs = {
             'input_ids': [],
@@ -19,101 +16,70 @@ def load_and_preprocess_data():
             'labels': []
         }
         for example in examples['prompt']:
-            # Tokenize chosen and rejected parts separately
-            chosen_tokenized = tokenizer(
+            # Tokenize input
+            tokenized = tokenizer(
                 example, 
                 truncation=True, 
                 padding='max_length', 
-                max_length=32
+                max_length=128  # Increase max_length to suit model needs
             )
-            rejected_tokenized = tokenizer(
-                example, 
-                truncation=True, 
-                padding='max_length', 
-                max_length=32
-            )
-            
-            combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
-            combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
-    
-            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)
-            inputs['labels'].append(torch.tensor([1] * (label_length // 16) + [0] * (label_length - label_length // 16)))
-    
-        # Convert lists to tensors for each batch
+
+            inputs['input_ids'].append(tokenized['input_ids'])
+            inputs['attention_mask'].append(tokenized['attention_mask'])
+            inputs['labels'].append(1 if "chosen" in example else 0)  # Assign binary labels
+
         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}")
-    
+        inputs['labels'] = torch.tensor(inputs['labels'], dtype=torch.long)
+
         return inputs
-    # Apply preprocessing to the dataset
-    dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
 
-    # Convert dataset to PyTorch format
+    dataset = dataset.map(preprocess_function, batched=True)
     dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
 
     return dataset, tokenizer
 
 # Define and train the model
 def train_model(dataset):
-    # Load the Longformer model for sequence classification
     model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2)
 
-    # Define training arguments
     training_args = TrainingArguments(
-        output_dir='./results',          
-        eval_strategy="epoch",            # Updated from evaluation_strategy
-        learning_rate=2e-5,              
-        per_device_train_batch_size=4,   
-        per_device_eval_batch_size=4,    
-        num_train_epochs=3,              
-        weight_decay=0.01,               
+        output_dir='./results',
+        evaluation_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        num_train_epochs=3,
+        weight_decay=0.01,
     )
 
-    # Initialize the Trainer with custom compute_loss to verify shapes
     class CustomTrainer(Trainer):
-        def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        def compute_loss(self, model, inputs, return_outputs=False):
+            labels = inputs.pop("labels")
             outputs = model(**inputs)
             logits = outputs.logits
-            labels = inputs['labels']
-            print(f"Logits shape: {logits.shape}, Labels shape: {labels.shape}")  # Debugging
-            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            loss = torch.nn.CrossEntropyLoss()(logits.view(-1, model.config.num_labels), labels.view(-1))
             return (loss, outputs) if return_outputs else loss
 
-    # Initialize the Trainer
     trainer = CustomTrainer(
         model=model,
         args=training_args,
         train_dataset=dataset,
-        eval_dataset=dataset,  # You can split the dataset for evaluation if needed
+        eval_dataset=dataset,
     )
 
-    # Train the model
     trainer.train()
-
     return model
 
 # Gradio Interface function to predict
 def predict(text):
-    # Tokenize the input text
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=32)  # Adjust max_length to match preprocessing
-    # Perform inference
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
     outputs = model(**inputs)
-    prediction = torch.argmax(outputs.logits, dim=-1).item()  # Get the class with the highest score
+    prediction = torch.argmax(outputs.logits, dim=-1).item()
     return f"Prediction: {'chosen' if prediction == 1 else 'rejected'}"
 
-# Main script to run training
+# Main script
 if __name__ == "__main__":
-    # Load and preprocess the data
     dataset, tokenizer = load_and_preprocess_data()
-
-    # Train the model
     model = train_model(dataset)
-
-    # Create a Gradio interface for prediction
-    gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()
+    gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()