Upload code/train_sft.py with huggingface_hub

code/train_sft.py

@@ -563,21 +563,24 @@ def main():
 
     root_dir = os.path.abspath(os.path.join(cur_dir, ".."))
     data_dir = args.data_dir or os.path.join(root_dir, "tokenized_data")
-    stage_subfolder = "ViuTranslate" if str(args.version).lower() == "viutranslate" else f"sft_{args.version}"
+    v_clean = str(args.version).lower().strip()
+    if v_clean in ("500m_sarus_sft", "sarus_sft", "500m"):
+        stage_subfolder = "500M_Sarus_SFT"
+        ckpt_tag = "500m_sarus_sft"
+    elif v_clean == "viutranslate":
+        stage_subfolder = "ViuTranslate"
+        ckpt_tag = "viutranslate"
+    else:
+        stage_subfolder = f"sft_{args.version}"
+        ckpt_tag = f"sft_{args.version}"
     stage_data_dir = os.path.join(data_dir, stage_subfolder)
 
-    output_dir = args.output_dir or (
-        os.path.join(root_dir, "sft_checkpoints", "ViuTranslate")
-        if str(args.version).lower() == "viutranslate"
-        else os.path.join(root_dir, "sft_checkpoints", f"sft_{args.version}")
-    )
+    output_dir = args.output_dir or os.path.join(root_dir, "sft_checkpoints", stage_subfolder)
     os.makedirs(output_dir, exist_ok=True)
     
-    ckpt_filename = "viutranslate_final.pt" if str(args.version).lower() == "viutranslate" else f"sft_{args.version}_final.pt"
+    ckpt_filename = f"{ckpt_tag}_final.pt"
     save_path = os.path.join(output_dir, ckpt_filename)
-    # Best-checkpoint is kept in a SEPARATE file so later (worse) evals can never
-    # overwrite the best weights. Previously both shared save_path.
-    best_filename = "viutranslate_best.pt" if str(args.version).lower() == "viutranslate" else f"sft_{args.version}_best.pt"
+    best_filename = f"{ckpt_tag}_best.pt"
     best_save_path = os.path.join(output_dir, best_filename)
     
     if args.init_ckpt: