changes to training script

scripts/train.py

@@ -102,7 +102,6 @@ def main():
         ):
             param.requires_grad = True
             trainable_params.append(name)
-            print(f"Unfreezing parameter: {name}")
 
     
     print(f"Total trainable parameters: {len(trainable_params)}")
@@ -192,13 +191,21 @@ def main():
     model.zero_grad()
 
     # Check for existing checkpoint
+    import re
+
     checkpoint_dir = None
     if os.path.isdir(training_args.output_dir):
-        checkpoints = [os.path.join(training_args.output_dir, d) for d in os.listdir(training_args.output_dir) if d.startswith("checkpoint-")]
+        checkpoints = [
+            os.path.join(training_args.output_dir, d)
+            for d in os.listdir(training_args.output_dir)
+            if re.match(r"checkpoint-\d+", d)
+        ]
         if checkpoints:
-            checkpoint_dir = max(checkpoints, key=os.path.getmtime)
+            # Extract step numbers and find the highest
+            checkpoint_dir = max(checkpoints, key=lambda x: int(x.split('-')[-1]))
             print(f"Resuming from checkpoint: {checkpoint_dir}")
 
+
     # Train
     print("Starting training...")
     trainer.train(resume_from_checkpoint=checkpoint_dir)