move details script and fix wandb logging (#314)

huggingface · Feb 13, 2025 · 80e7e7b · 80e7e7b
1 parent f987b3c
commit 80e7e7b
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 10 deletions.
diff --git a/src/open_r1/utils/upload_details.py → scripts/upload_details.py b/src/open_r1/utils/upload_details.py → scripts/upload_details.py
diff --git a/slurm/evaluate.slurm b/slurm/evaluate.slurm
@@ -81,7 +81,7 @@ echo "Uploading details to Hugging Face Hub..."
 DETAILS_FILEPATHS=$(find $OUTPUT_DIR/details/ -type f \( -name "*.parquet" \))
 echo "DETAILS_FILEPATHS: $DETAILS_FILEPATHS"
 TIMESTAMP=$(date +"%Y-%m-%dT%H-%M-%S")
-python src/open_r1/utils/upload_details.py --data_files $DETAILS_FILEPATHS --hub_repo_id $DETAILS_REPO_ID --config_name $MODEL_REVISION.$TASK_NAME.$TIMESTAMP
+python scripts/upload_details.py --data_files $DETAILS_FILEPATHS --hub_repo_id $DETAILS_REPO_ID --config_name $MODEL_REVISION.$TASK_NAME.$TIMESTAMP
 
 echo "Cleaning up ..."
 rm -rf $OUTPUT_DIR

diff --git a/src/open_r1/utils/logging.py b/src/open_r1/utils/logging.py
diff --git a/src/open_r1/utils/wandb_logging.py b/src/open_r1/utils/wandb_logging.py
@@ -0,0 +1,11 @@
+import os
+
+
+def init_wandb_training(training_args):
+    """
+    Helper function for setting up Weights & Biases logging tools.
+    """
+    if training_args.wandb_entity is not None:
+        os.environ["WANDB_ENTITY"] = training_args.wandb_entity
+    if training_args.wandb_project is not None:
+        os.environ["WANDB_PROJECT"] = training_args.wandb_project