Update train_model.py

train_model.py

@@ -67,8 +67,20 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
         logging.info("Dataset loaded successfully.")
 
         def tokenize_function(examples):
-            # Truncate and set max_length, but let DataCollator handle padding
-            return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
+            try:
+                # Tokenize with truncation, defer padding to DataCollator
+                tokens = tokenizer(
+                    examples['text'],
+                    truncation=True,
+                    max_length=sequence_length,  # Set maximum length
+                    padding=False,  # Padding will be handled by the DataCollatorWithPadding
+                    return_tensors=None  # Let the DataCollator handle tensor creation
+                )
+                return tokens
+            except Exception as e:
+                logging.error(f"Error during tokenization: {e}")
+                logging.error(f"Example data: {examples}")
+                raise e
 
         # Tokenize the dataset using the modified tokenize_function
         tokenized_datasets = dataset.shuffle(seed=42).select(range(500)).map(tokenize_function, batched=True)
@@ -182,7 +194,6 @@ 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(
@@ -194,26 +205,38 @@ 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, padding=True)  # Explicit padding
+        data_collator = DataCollatorWithPadding(tokenizer=tokenizer)  # Dynamic padding during batch creation
     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 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)  # None for scheduler
+        optimizers=(get_optimizer(model, training_args.learning_rate), None)
     )
 
-
     # Start training
     logging.info("Starting training...")
     try:
@@ -253,5 +276,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-