dev

app.py

@@ -241,17 +241,14 @@ def fine_tune_model(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, g
         def get_checkpoint_int(s):
             int_index = s.find('-')
             return int(s[int_index+1:])
-            
-        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/checkpoints"), reverse=True)
-        print(f'CHECKPOINTs: {previous_checkpoints}')
-        for check in previous_checkpoints:
-            try:
-                print(f"Removing previous checkpoint {check}")
-                shutil.rmtree(os.path.join('/data/results/checkpoints', check)) 
-            except:
-                pass        
+                
+        def filter_checkpoints_dirs(l):
+            new_list = list()
+            for entry in l:
+                if 'checkpoint' in entry:
+                    new_list.append(entry)
+            return new_list
+                
         try:
             train_result = trainer.train(resume_from_checkpoint=True)
         except Exception as e:
@@ -259,12 +256,13 @@ def fine_tune_model(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, g
             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)
+            dir_entries = filter_checkpoints_dirs(os.listdir(checkpoint_dir))
+            previous_checkpoints = sorted(dir_entries, key=get_checkpoint_int, reverse=True)
             print(f'CHECKPOINTs: {previous_checkpoints}')
             for check in previous_checkpoints:
                 try:
                     print(f"Removing previous checkpoint {check}")
-                    shutil.rmtree(os.path.join('/data/results', check))             
+                    shutil.rmtree(os.path.join(checkpoint_dir, check))             
                     train_result = trainer.train(resume_from_checkpoint=True)
                     trainer.push_to_hub(commit_message="Training complete!")
                     return 'DONE!'#train_result