dev

app.py

@@ -231,29 +231,32 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             ##compute_metrics=compute_metrics,
             ##data_collator=data_collator,
             ##processing_class=tokenizer,
-        #)    
-        checkpoint_dir = training_args.output_dir
-        if os.path.exists(checkpoint_dir) and os.listdir(checkpoint_dir):
-            # Check if the trainer_state.json file exists in the specified checkpoint
-            trainer_state_path = os.path.join(checkpoint_dir, 'trainer_state.json')
-            if os.path.exists(trainer_state_path):
-                train_result = trainer.train(resume_from_checkpoint=True)
-            else:
-                # If the trainer_state.json is missing, look for the previous checkpoint
-                print(f"Checkpoint {checkpoint_dir} is missing 'trainer_state.json'. Looking for previous checkpoints...")
-                previous_checkpoints = sorted(glob.glob(os.path.join(os.path.dirname(checkpoint_dir), 'checkpoint-*')), key=os.path.getmtime)
-                
-                if previous_checkpoints:
-                    # Load the most recent previous checkpoint
-                    last_checkpoint = previous_checkpoints[-1]
-                    print(f"Loading previous checkpoint: {last_checkpoint}")
-                    train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
+        #)
+        try:
+            train_result = trainer.train(resume_from_checkpoint=True)
+        except:
+            checkpoint_dir = training_args.output_dir
+            if os.path.exists(checkpoint_dir) and os.listdir(checkpoint_dir):
+                # Check if the trainer_state.json file exists in the specified checkpoint
+                trainer_state_path = os.path.join(checkpoint_dir, 'trainer_state.json')
+                if os.path.exists(trainer_state_path):
+                    train_result = trainer.train(resume_from_checkpoint=True)
                 else:
-                    print("No previous checkpoints found. Starting training from scratch.")
-                    train_result = trainer.train()
-        else:
-            print("No checkpoints found. Starting training from scratch.")
-            train_result = trainer.train()
+                    # If the trainer_state.json is missing, look for the previous checkpoint
+                    print(f"Checkpoint {checkpoint_dir} is missing 'trainer_state.json'. Looking for previous checkpoints...")
+                    previous_checkpoints = sorted(glob.glob(os.path.join(os.path.dirname(checkpoint_dir), 'checkpoint-*')), key=os.path.getmtime)
+                    
+                    if previous_checkpoints:
+                        # Load the most recent previous checkpoint
+                        last_checkpoint = previous_checkpoints[-1]
+                        print(f"Loading previous checkpoint: {last_checkpoint}")
+                        train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
+                    else:
+                        print("No previous checkpoints found. Starting training from scratch.")
+                        train_result = trainer.train()
+            else:
+                print("No checkpoints found. Starting training from scratch.")
+                train_result = trainer.train()
 
         trainer.push_to_hub(commit_message="Training complete!")
     except Exception as e: