Add comments to generated config

yourbench_space/config.py

@@ -1,7 +1,8 @@
-import yaml
+from ruamel.yaml import YAML
 from loguru import logger
 
 from yourbench_space import PATH
+from yourbench_space.utils import to_commentable_yaml
 
 
 def generate_base_config(hf_org: str, hf_dataset_name: str, session_uid: str):
@@ -82,10 +83,24 @@ def generate_base_config(hf_org: str, hf_dataset_name: str, session_uid: str):
     }
 
 
-def save_yaml_file(config: str, path: str):
-    """Saves the given config dictionary to a YAML file"""
+def save_yaml_file(config: dict, path: str):
+    """Saves the given config dictionary to a YAML file with helpful comments."""
+    yaml = YAML()
+    yaml.indent(mapping=2, sequence=4, offset=2)
+
+    config_cm = to_commentable_yaml(config)
+
+    # Now we can add inline comments
+    ingestion = config_cm["pipeline"]["ingestion"]
+    ingestion.yaml_set_comment_before_after_key("source_documents_dir", before="⚠️ Change this path to match your local directory")
+    ingestion.yaml_set_comment_before_after_key("output_dir", before="⚠️ This is where ingested data will be saved")
+
+    upload = config_cm["pipeline"]["upload_ingest_to_hub"]
+    upload.yaml_set_comment_before_after_key("source_documents_dir", before="⚠️ Same as output_dir from ingestion — adjust as needed")
+
     with open(path, "w") as file:
-        yaml.dump(config, file, default_flow_style=False, sort_keys=False)
+        yaml.dump(config_cm, file)
+
     return path
 
 

yourbench_space/utils.py

@@ -5,6 +5,7 @@ import shutil
 import pathlib
 import subprocess
 from typing import List, Union, Optional
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 
 import pandas as pd
 from loguru import logger
@@ -34,6 +35,22 @@ STAGE_DISPLAY_MAP = {
     "lighteval": "Generate Lighteval Subset",
 }
 
+def to_commentable_yaml(obj):
+    """
+    Recursively converts standard Python dicts and lists into
+    ruamel.yaml's CommentedMap and CommentedSeq so that comments
+    can be attached when dumping YAML
+    """
+    # Convert dict to CommentedMap with recursively processed values
+    if isinstance(obj, dict):
+        return CommentedMap({k: to_commentable_yaml(v) for k, v in obj.items()})
+    
+    # Convert list to CommentedSeq with recursively processed elements
+    elif isinstance(obj, list):
+        return CommentedSeq([to_commentable_yaml(i) for i in obj])
+    
+    # Return non-container values as-is
+    return obj
 
 def map_stage_names(stages: list[str]) -> list[str]:
     return [STAGE_DISPLAY_MAP.get(stage, stage) for stage in stages]