Upload code/train_dpo.py with huggingface_hub

code/train_dpo.py

@@ -177,8 +177,12 @@ def load_dpo_data():
 
     if not os.path.exists(local_train) and IS_MAIN:
         print("  ⬇️ Fetching DPO pairs from Hugging Face Dataset Hub...")
-        dl_t = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_train.json", repo_type="dataset", token=HF_TOKEN)
-        dl_v = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_val.json", repo_type="dataset", token=HF_TOKEN)
+        try:
+            dl_t = hf_hub_download(repo_id="ViuAI/500M_Sarus_DPO", filename="dpo_train.json", repo_type="dataset", token=HF_TOKEN)
+            dl_v = hf_hub_download(repo_id="ViuAI/500M_Sarus_DPO", filename="dpo_val.json", repo_type="dataset", token=HF_TOKEN)
+        except Exception:
+            dl_t = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_train.json", repo_type="dataset", token=HF_TOKEN)
+            dl_v = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_val.json", repo_type="dataset", token=HF_TOKEN)
         os.makedirs(os.path.dirname(local_train), exist_ok=True)
         shutil.copy(dl_t, local_train)
         shutil.copy(dl_v, local_val)