Update train_model.py

train_model.py

@@ -70,9 +70,9 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
             # 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')
+                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train', use_auth_token=True)
             else:
-                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
+                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train', use_auth_token=True)
             logging.info("Dataset loaded successfully for generation task.")
             def tokenize_function(examples):
                 return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
@@ -185,6 +185,8 @@ 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}")
+            # Resize model's token embeddings after setting pad_token
             model = initialize_model(
                 task=args.task,
                 model_name=args.model_name,
@@ -195,7 +197,10 @@ def main():
                 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,
@@ -206,7 +211,7 @@ def main():
                 attention_heads=args.attention_heads
             )
     except Exception as e:
-        logging.error(f"Error initializing tokenizer: {str(e)}")
+        logging.error(f"Error initializing tokenizer or model: {str(e)}")
         raise e
 
     # Load and prepare dataset
@@ -221,9 +226,6 @@ def main():
         logging.error("Failed to load and prepare dataset.")
         raise e
 
-    # Initialize model (Already initialized above)
-    # model = initialize_model(...)  # Moved above to handle pad_token
-
     # Define data collator
     if args.task == "generation":
         data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
@@ -245,7 +247,7 @@ def main():
             learning_rate=5e-4,
             remove_unused_columns=False,
             push_to_hub=False  # We'll handle pushing manually
-           
+            
         )
     elif args.task == "classification":
         training_args = TrainingArguments(
@@ -313,8 +315,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
-
-
-