dev

app.py

@@ -1,4 +1,5 @@
 import spaces
+import glob
 import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import DataCollatorForSeq2Seq, AutoConfig
@@ -231,12 +232,29 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             ##data_collator=data_collator,
             ##processing_class=tokenizer,
         #)    
-    
-        # Fine-tune the model
-        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
-            train_result = trainer.train(resume_from_checkpoint=True)
+        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)
+                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:
         return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"