init

app.py

@@ -59,7 +59,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=32)
             
             tokenized_datasets['train'].save_to_disk(f'{hub_id.strip()}_train_dataset')
-            tokenized_datasets['validation'].save_to_disk(f'{hub_id.strip()}_test_dataset')
+            tokenized_datasets['test'].save_to_disk(f'{hub_id.strip()}_test_dataset')
         
     
         # Set training arguments
@@ -94,8 +94,8 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=data_collator['train'],
-            eval_dataset=data_collator['test'],
+            train_dataset=tokenized_datasets['train'],
+            eval_dataset=tokenized_datasets['test'],
             #callbacks=[LoggingCallback()], 
         )