Update train_model.py

train_model.py

@@ -182,6 +182,7 @@ def main():
         logging.error(f"Error initializing tokenizer or model: {str(e)}")
         raise e
 
+    # Load and prepare dataset
     # Load and prepare dataset
     try:
         tokenized_datasets = load_and_prepare_dataset(
@@ -193,38 +194,26 @@ def main():
     except Exception as e:
         logging.error("Failed to load and prepare dataset.")
         raise e
-
+    
     # Define data collator
     if args.task == "generation":
         data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
     elif args.task == "classification":
-        data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+        data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=True)  # Explicit padding
     else:
         logging.error("Unsupported task type for data collator.")
         raise ValueError("Unsupported task type for data collator.")
-
-    # Define training arguments
-    training_args = TrainingArguments(
-        output_dir=f"./models/{args.model_name}",
-        num_train_epochs=3,
-        per_device_train_batch_size=8 if args.task == "generation" else 16,
-        save_steps=5000,
-        save_total_limit=2,
-        logging_steps=500,
-        learning_rate=5e-4 if args.task == "generation" else 5e-5,
-        remove_unused_columns=False,
-        push_to_hub=False
-    )
-
-    # Initialize Trainer with PyTorch's AdamW optimizer
+    
+    # Initialize Trainer with the data collator
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=tokenized_datasets,
         data_collator=data_collator,
-        optimizers=(get_optimizer(model, training_args.learning_rate), None)
+        optimizers=(get_optimizer(model, training_args.learning_rate), None)  # None for scheduler
     )
 
+
     # Start training
     logging.info("Starting training...")
     try: