Update train_model.py

train_model.py

@@ -16,7 +16,6 @@ import torch
 import os
 from huggingface_hub import login, HfApi
 import logging
-
 from torch.optim import AdamW  # Import PyTorch's AdamW
 
 def setup_logging(log_file_path):
@@ -64,18 +63,14 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
     """
     logging.info(f"Loading dataset '{dataset_name}' for task '{task}'...")
     try:
-        if '/' in dataset_name:
-            dataset, config = dataset_name.split('/', 1)
-            dataset = load_dataset("stanfordnlp/imdb",split='train')
-        else:
-            dataset = load_dataset("stanfordnlp/imdb",split='train')
-        
+        dataset = load_dataset(dataset_name, split='train')
         logging.info("Dataset loaded successfully.")
 
         def tokenize_function(examples):
-            return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=sequence_length)
-        
-        # Tokenize the dataset
+            # Truncate and set max_length, but let DataCollator handle padding
+            return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
+
+        # Tokenize the dataset using the modified tokenize_function
         tokenized_datasets = dataset.shuffle(seed=42).select(range(500)).map(tokenize_function, batched=True)
         logging.info("Dataset tokenization complete.")
         return tokenized_datasets
@@ -100,7 +95,6 @@ def initialize_model(task, model_name, vocab_size, sequence_length, hidden_size,
                 intermediate_size=4 * hidden_size,
                 hidden_act='gelu',
                 use_cache=True,
-                truncation=False
             )
             model = GPT2LMHeadModel(config)
             logging.info("GPT2LMHeadModel initialized successfully.")
@@ -172,31 +166,18 @@ def main():
         if tokenizer.pad_token is None:
             logging.info("Setting pad_token to eos_token.")
             tokenizer.pad_token = tokenizer.eos_token
-            logging.info(f"Tokenizer pad_token set to: {tokenizer.pad_token}")
-            # Initialize model after setting pad_token
-            model = initialize_model(
-                task=args.task,
-                model_name=args.model_name,
-                vocab_size=args.vocab_size,
-                sequence_length=args.sequence_length,
-                hidden_size=args.hidden_size,
-                num_layers=args.num_layers,
-                attention_heads=args.attention_heads
-            )
-            model.resize_token_embeddings(len(tokenizer))
-            logging.info("Resized token embeddings to accommodate pad_token.")
-        else:
-            logging.info(f"Tokenizer already has pad_token set to: {tokenizer.pad_token}")
-            # Initialize model normally
-            model = initialize_model(
-                task=args.task,
-                model_name=args.model_name,
-                vocab_size=args.vocab_size,
-                sequence_length=args.sequence_length,
-                hidden_size=args.hidden_size,
-                num_layers=args.num_layers,
-                attention_heads=args.attention_heads
-            )
+        
+        # Initialize model
+        model = initialize_model(
+            task=args.task,
+            model_name=args.model_name,
+            vocab_size=args.vocab_size,
+            sequence_length=args.sequence_length,
+            hidden_size=args.hidden_size,
+            num_layers=args.num_layers,
+            attention_heads=args.attention_heads
+        )
+        model.resize_token_embeddings(len(tokenizer))
     except Exception as e:
         logging.error(f"Error initializing tokenizer or model: {str(e)}")
         raise e
@@ -223,36 +204,17 @@ def main():
         raise ValueError("Unsupported task type for data collator.")
 
     # Define training arguments
-    if args.task == "generation":
-        training_args = TrainingArguments(
-            output_dir=f"./models/{args.model_name}",
-            num_train_epochs=3,
-            per_device_train_batch_size=8,
-            save_steps=5000,
-            save_total_limit=2,
-            logging_steps=500,
-            learning_rate=5e-4,
-            remove_unused_columns=False,
-            push_to_hub=False  # We'll handle pushing manually
-            
-        )
-    elif args.task == "classification":
-        training_args = TrainingArguments(
-            output_dir=f"./models/{args.model_name}",
-            num_train_epochs=3,
-            per_device_train_batch_size=16,
-            evaluation_strategy="epoch",
-            save_steps=5000,
-            save_total_limit=2,
-            logging_steps=500,
-            learning_rate=5e-5,
-            remove_unused_columns=False,
-            push_to_hub=False  # We'll handle pushing manually
-            
-        )
-    else:
-        logging.error("Unsupported task type for training arguments.")
-        raise ValueError("Unsupported task type for 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
     trainer = Trainer(
@@ -260,7 +222,7 @@ def main():
         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
@@ -303,3 +265,4 @@ def main():
 if __name__ == "__main__":
     main()
 
+