Upload code/train_sft.py with huggingface_hub

code/train_sft.py

@@ -325,13 +325,13 @@ def main():
     stage_subfolder = f"sft_{args.version}"
     stage_data_dir = os.path.join(data_dir, stage_subfolder)
     
-    final_save_path = os.path.join(output_dir, f"sft_{args.version}_final.pt")
-    best_save_path = os.path.join(output_dir, f"sft_{args.version}_best.pt")
+    ckpt_filename = f"sft_{args.version}_final.pt"
+    save_path = os.path.join(output_dir, ckpt_filename)
     
     if args.init_ckpt:
         init_ckpt = args.init_ckpt
-    elif args.resume and os.path.exists(best_save_path):
-        init_ckpt = best_save_path
+    elif args.resume and os.path.exists(save_path):
+        init_ckpt = save_path
     else:
         init_ckpt = os.path.join(root_dir, "checkpoints", "ckpt_latest.pt")
 
@@ -653,12 +653,8 @@ def main():
                         "model_args": vars(model_args),
                         "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
                     }
-                    # Save both step checkpoint and best checkpoint
-                    if is_best:
-                        torch.save(save_payload, best_save_path)
-                        print(f"  💾 Saved BEST checkpoint -> {best_save_path}")
-                    torch.save(save_payload, final_save_path)
-                    print(f"  💾 Saved LATEST checkpoint -> {final_save_path}\n")
+                    torch.save(save_payload, save_path)
+                    print(f"  💾 Saved checkpoint -> {save_path}\n")
 
         # End of Epoch Handling
         if (micro_idx + 1) % grad_acc != 0:
@@ -696,32 +692,24 @@ def main():
         "model_args": vars(model_args),
         "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
     }
-    torch.save(final_payload, final_save_path)
-    print(f"💾 Final master checkpoint saved to: {final_save_path}")
+    torch.save(final_payload, save_path)
+    print(f"💾 Final master checkpoint saved to: {save_path}")
 
-    # Push to Hugging Face Hub (Both Best & Final Checkpoints)
+    # Push to Hugging Face Hub (Only Final Checkpoint)
     if args.push_to_hf:
-        print("\n🚀 Pushing Best & Final Checkpoints to Hugging Face Model Hub (ViuAI/ViuAI-500M)...")
+        print("\n🚀 Pushing Final Checkpoint to Hugging Face Model Hub (ViuAI/ViuAI-500M)...")
         token = args.hf_token or os.environ.get("HF_TOKEN")
         if token:
             try:
                 api = HfApi(token=token)
-                if os.path.exists(best_save_path):
-                    api.upload_file(
-                        path_or_fileobj=best_save_path,
-                        path_in_repo=f"checkpoints/sft_{args.version}_best.pt",
-                        repo_id="ViuAI/ViuAI-500M",
-                        repo_type="model"
-                    )
-                    print(f"✅ Successfully uploaded sft_{args.version}_best.pt!")
-                if os.path.exists(final_save_path):
+                if os.path.exists(save_path):
                     api.upload_file(
-                        path_or_fileobj=final_save_path,
-                        path_in_repo=f"checkpoints/sft_{args.version}_final.pt",
+                        path_or_fileobj=save_path,
+                        path_in_repo=f"checkpoints/{ckpt_filename}",
                         repo_id="ViuAI/ViuAI-500M",
                         repo_type="model"
                     )
-                    print(f"✅ Successfully uploaded sft_{args.version}_final.pt!")
+                    print(f"✅ Successfully uploaded {ckpt_filename} to ViuAI/ViuAI-500M!")
             except Exception as e:
                 print(f"⚠️ Error uploading to Hugging Face: {e}")
         else: