Update app.py

app.py

@@ -13,7 +13,6 @@ def load_and_preprocess_data():
 
     # 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': [],
@@ -36,7 +35,7 @@ def load_and_preprocess_data():
             
             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]
@@ -45,14 +44,20 @@ 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'].append(combined_input_ids)
             inputs['attention_mask'].append(combined_attention_mask)
             inputs['labels'].append([1] * 500 + [0] * 500)  # Adjust based on your label structure
-    
+
+        # 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.tensor(inputs['labels'])
+
         return inputs
+
     # Apply preprocessing to the dataset
-    dataset = dataset.map(preprocess_function, batched=True)
+    dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
 
     # Convert dataset to PyTorch format
     dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
@@ -66,17 +71,27 @@ def train_model(dataset):
 
     # Define training arguments
     training_args = TrainingArguments(
-        output_dir='./results',          # Output directory for model checkpoints
-        evaluation_strategy="epoch",     # Evaluation at the end of each epoch
-        learning_rate=2e-5,              # Learning rate
-        per_device_train_batch_size=4,   # Reduced batch size due to longer sequences
-        per_device_eval_batch_size=4,    # Reduced batch size for evaluation
-        num_train_epochs=3,              # Number of training epochs
-        weight_decay=0.01,               # Strength of weight decay
+        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,               
     )
 
+    # Initialize the Trainer with custom compute_loss to verify shapes
+    class CustomTrainer(Trainer):
+        def compute_loss(self, model, inputs, return_outputs=False):
+            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))
+            return (loss, outputs) if return_outputs else loss
+
     # Initialize the Trainer
-    trainer = Trainer(
+    trainer = CustomTrainer(
         model=model,
         args=training_args,
         train_dataset=dataset,
@@ -91,7 +106,7 @@ def train_model(dataset):
 # Gradio Interface function to predict
 def predict(text):
     # Tokenize the input text
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=4096)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=1000)  # Adjust max_length to match preprocessing
     # Perform inference
     outputs = model(**inputs)
     prediction = torch.argmax(outputs.logits, dim=-1).item()  # Get the class with the highest score
@@ -105,5 +120,4 @@ if __name__ == "__main__":
     # Train the model
     model = train_model(dataset)
 
-    # Create a Gradio interface for prediction
-    gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()
+    # Create