Update train_model.py

train_model.py

@@ -12,12 +12,14 @@ from transformers import (
     DataCollatorForLanguageModeling,
     DataCollatorWithPadding,
 )
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
 import torch
 import os
 from huggingface_hub import login, HfApi, HfFolder
 import logging
 
+from torch.optim import AdamW  # Import PyTorch's AdamW
+
 def setup_logging(log_file_path):
     """
     Sets up logging to both console and a file.
@@ -65,21 +67,21 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
     logging.info(f"Loading dataset '{dataset_name}' for task '{task}'...")
     try:
         if task == "generation":
-            # Check if dataset_name includes config
+            # Check if dataset_name includes a configuration
             if '/' in dataset_name:
                 dataset, config = dataset_name.split('/', 1)
-                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train[:1%]')
+                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
             else:
-                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train[:1%]')
+                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
             logging.info("Dataset loaded successfully for generation task.")
             def tokenize_function(examples):
                 return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
         elif task == "classification":
             if '/' in dataset_name:
                 dataset, config = dataset_name.split('/', 1)
-                dataset = load_dataset("stanfordnlp/imdb", split='train[:1%]')
+                dataset = load_dataset(dataset, config, split='train', use_auth_token=True)
             else:
-                dataset = load_dataset("stanfordnlp/imdb", split='train[:1%]')
+                dataset = load_dataset(dataset_name, split='train', use_auth_token=True)
             logging.info("Dataset loaded successfully for classification task.")
             # Assuming the dataset has 'text' and 'label' columns
             def tokenize_function(examples):
@@ -136,6 +138,12 @@ def initialize_model(task, model_name, vocab_size, sequence_length, hidden_size,
         logging.error(f"Error initializing model: {str(e)}")
         raise e
 
+def get_optimizer(model, learning_rate):
+    """
+    Returns the AdamW optimizer from PyTorch.
+    """
+    return AdamW(model.parameters(), lr=learning_rate)
+
 def main():
     # Parse arguments
     args = parse_arguments()
@@ -172,6 +180,31 @@ def main():
         else:
             raise ValueError("Unsupported task type")
         logging.info("Tokenizer initialized successfully.")
+        
+        # Set pad_token to eos_token if not already set
+        if tokenizer.pad_token is None:
+            logging.info("Setting pad_token to eos_token.")
+            tokenizer.pad_token = tokenizer.eos_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))
+        else:
+            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
+            )
     except Exception as e:
         logging.error(f"Error initializing tokenizer: {str(e)}")
         raise e
@@ -188,20 +221,8 @@ def main():
         logging.error("Failed to load and prepare dataset.")
         raise e
 
-    # Initialize model
-    try:
-        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
-        )
-    except Exception as e:
-        logging.error("Failed to initialize model.")
-        raise e
+    # Initialize model (Already initialized above)
+    # model = initialize_model(...)  # Moved above to handle pad_token
 
     # Define data collator
     if args.task == "generation":
@@ -223,7 +244,8 @@ def main():
             logging_steps=500,
             learning_rate=5e-4,
             remove_unused_columns=False,
-            push_to_hub=False  # We'll handle pushing manually
+            push_to_hub=False,  # We'll handle pushing manually
+            no_deprecation_warning=True  # Suppress FutureWarning
         )
     elif args.task == "classification":
         training_args = TrainingArguments(
@@ -236,18 +258,20 @@ def main():
             logging_steps=500,
             learning_rate=5e-5,
             remove_unused_columns=False,
-            push_to_hub=False  # We'll handle pushing manually
+            push_to_hub=False,  # We'll handle pushing manually
+            no_deprecation_warning=True  # Suppress FutureWarning
         )
     else:
         logging.error("Unsupported task type for training arguments.")
         raise ValueError("Unsupported task type for training arguments.")
 
-    # Initialize Trainer
+    # Initialize Trainer with PyTorch's AdamW optimizer
     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
     )
 
     # Start training
@@ -293,3 +317,4 @@ if __name__ == "__main__":
 
 
 
+