Update train_model.py

train_model.py

@@ -65,13 +65,13 @@ 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":
-            train_dataset = load_dataset(dataset_name,use_auth_token=True)
+            train_dataset = load_dataset(dataset_name,split='train',use_auth_token=True)
             dataset = train_dataset['train'].shuffle(seed=42).select(range(500))
             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":
-            train_dataset = load_dataset(dataset_name,use_auth_token=True)
+            train_dataset = load_dataset(dataset_name,split='train',use_auth_token=True)
             dataset = train_dataset['train'].shuffle(seed=42).select(range(500))
             logging.info("Dataset loaded successfully for classification task.")
             # Assuming the dataset has 'text' and 'label' columns