Update train_model.py

train_model.py

@@ -16,7 +16,7 @@ import torch
 import os
 from huggingface_hub import login, HfApi
 import logging
-from torch.optim import AdamW  # Import PyTorch's AdamW
+from torch.optim import AdamW
 
 def setup_logging(log_file_path):
     """
@@ -66,23 +66,28 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
         dataset = load_dataset(dataset_name, split='train')
         logging.info("Dataset loaded successfully.")
 
+        # Log some examples to check dataset structure
+        logging.info(f"Example data from the dataset: {dataset[:5]}")
+
         def tokenize_function(examples):
             try:
-                # Tokenize with truncation, defer padding to DataCollator
+                # Tokenize with truncation and padding
                 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
+                    max_length=sequence_length,
+                    padding='max_length',  # Force padding to max length for debugging
+                    return_tensors=None  # Let the collator handle tensor conversion
                 )
+                # Log the tokens for debugging
+                logging.info(f"Tokenized example: {tokens}")
                 return tokens
             except Exception as e:
                 logging.error(f"Error during tokenization: {e}")
-                logging.error(f"Example data: {examples}")
+                logging.error(f"Problematic example: {examples}")
                 raise e
 
-        # Tokenize the dataset using the modified tokenize_function
+        # Tokenize the dataset
         tokenized_datasets = dataset.shuffle(seed=42).select(range(500)).map(tokenize_function, batched=True)
         logging.info("Dataset tokenization complete.")
         return tokenized_datasets
@@ -210,7 +215,7 @@ def main():
     if args.task == "generation":
         data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
     elif args.task == "classification":
-        data_collator = DataCollatorWithPadding(tokenizer=tokenizer)  # Dynamic padding during batch creation
+        data_collator = DataCollatorWithPadding(tokenizer=tokenizer)  # Handle padding dynamically during batching
     else:
         logging.error("Unsupported task type for data collator.")
         raise ValueError("Unsupported task type for data collator.")
@@ -276,3 +281,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+