dev

app.py

@@ -244,6 +244,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         try:
             train_result = trainer.train(resume_from_checkpoint=True)
         except:
+            import shutil
             checkpoint_dir = training_args.output_dir
             # If the trainer_state.json is missing, look for the previous checkpoint
             previous_checkpoints = sorted(os.listdir("/data/results"), key=get_checkpoint_int, reverse=True)
@@ -251,7 +252,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             for check in previous_checkpoints:
                 try:
                     print(f"Removing previous checkpoint {check}")
-                    os.remove(os.path.join('/data/results', check))                    
+                    shutil.rmtree(os.path.join('/data/results', check))             
                     train_result = trainer.train(resume_from_checkpoint=True)
                     trainer.push_to_hub(commit_message="Training complete!")
                     return 'DONE!'#train_result