Update train_model.py

train_model.py

@@ -12,10 +12,10 @@ from transformers import (
     DataCollatorForLanguageModeling,
     DataCollatorWithPadding,
 )
-from datasets import load_dataset, Dataset
+from datasets import load_dataset
 import torch
 import os
-from huggingface_hub import login, HfApi, HfFolder
+from huggingface_hub import login, HfApi
 import logging
 
 from torch.optim import AdamW  # Import PyTorch's AdamW
@@ -34,10 +34,9 @@ def setup_logging(log_file_path):
     f_handler.setLevel(logging.INFO)
 
     # Create formatters and add to handlers
-    c_format = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-    f_format = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-    c_handler.setFormatter(c_format)
-    f_handler.setFormatter(f_format)
+    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+    c_handler.setFormatter(formatter)
+    f_handler.setFormatter(formatter)
 
     # Add handlers to the logger
     logger.addHandler(c_handler)
@@ -66,30 +65,18 @@ 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 a configuration
-            if '/' in dataset_name:
-                dataset, config = dataset_name.split('/', 1)
-                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
-            else:
-                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(dataset, config, split='train')
-            else:
-                dataset = load_dataset(dataset_name, split='train')
-            logging.info("Dataset loaded successfully for classification task.")
-            # Assuming the dataset has 'text' and 'label' columns
-            def tokenize_function(examples):
-                return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
+        if '/' in dataset_name:
+            dataset, config = dataset_name.split('/', 1)
+            dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
         else:
-            raise ValueError("Unsupported task type")
+            dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
         
-        # Shuffle and select a subset
+        logging.info("Dataset loaded successfully.")
+
+        def tokenize_function(examples):
+            return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
+        
+        # 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
@@ -186,7 +173,7 @@ def main():
             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}")
-            # Resize model's token embeddings after setting pad_token
+            # Initialize model after setting pad_token
             model = initialize_model(
                 task=args.task,
                 model_name=args.model_name,
@@ -315,3 +302,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+