Update train_model.py

train_model.py

@@ -70,18 +70,18 @@ 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', use_auth_token=True)
+                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
             else:
-                dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train', use_auth_token=True)
+                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', use_auth_token=True)
+                dataset = load_dataset(dataset, config, split='train')
             else:
-                dataset = load_dataset(dataset_name, split='train', use_auth_token=True)
+                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):