dev

app.py

@@ -28,8 +28,8 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
     
         max_length = 128
         try:
-            tokenized_train_dataset = load_from_disk(f'{hub_id.strip()}_train_dataset')
-            tokenized_test_dataset = load_from_disk(f'{hub_id.strip()}_test_dataset')
+            tokenized_train_dataset = load_from_disk(f'data/{hub_id.strip()}_train_dataset')
+            tokenized_test_dataset = load_from_disk(f'data/{hub_id.strip()}_test_dataset')
             tokenized_datasets = concatenate_datasets([tokenized_train_dataset, tokenized_test_dataset])
         except:
             # Tokenize the dataset
@@ -58,8 +58,8 @@ 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['test'].save_to_disk(f'{hub_id.strip()}_test_dataset')
+            tokenized_datasets['train'].save_to_disk(f'data/{hub_id.strip()}_train_dataset')
+            tokenized_datasets['test'].save_to_disk(f'data/{hub_id.strip()}_test_dataset')
         
     
         # Set training arguments
@@ -98,8 +98,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             eval_dataset=tokenized_datasets['test'],
             #callbacks=[LoggingCallback()], 
         )
-        for batch in trainer.get_train_dataloader():
-            print(batch['input_ids'].shape, batch['labels'].shape)    
+
         # Fine-tune the model
         trainer.train()
         trainer.push_to_hub(commit_message="Training complete!")