Add progress_format support to datasets.utils.tqdm

Similar to huggingface/tokenizers#1921, adds machine-readable JSON progress output.

- Add set_progress_format() and get_progress_format() functions
- Support 'tqdm' (default), 'json', and 'silent' formats
- Emit JSON progress every 5% when format='json'
- Export new functions from datasets.utils

Cross-reference: huggingface/tokenizers#1921

Author: podarok

src/datasets/utils/__init__.py

@@ -20,6 +20,8 @@
     are_progress_bars_disabled,
     disable_progress_bars,
     enable_progress_bars,
+    get_progress_format,
+    set_progress_format,
     tqdm,
 )
 from .version import Version

src/datasets/utils/tqdm.py

@@ -6,6 +6,7 @@
        environment variable `HF_DATASETS_DISABLE_PROGRESS_BARS` to 1.
     3. To re-enable progress bars, use `enable_progress_bars()`.
     4. To check whether progress bars are disabled, use `are_progress_bars_disabled()`.
+    5. To emit machine-readable JSON progress, use `set_progress_format("json")`.
 
 NOTE: Environment variable `HF_DATASETS_DISABLE_PROGRESS_BARS` has the priority.
 
@@ -15,6 +16,8 @@
         are_progress_bars_disabled,
         disable_progress_bars,
         enable_progress_bars,
+        set_progress_format,
+        get_progress_format,
         tqdm,
     )
 
@@ -37,9 +40,17 @@
     # Progress bar will be shown !
     for _ in tqdm(range(5)):
        do_something()
+
+    # Emit JSON progress (machine-readable)
+    set_progress_format("json")
+    for i in tqdm(range(100)):
+       do_something()
+    # Outputs: {"stage":"","current":50,"total":100,"percent":50.0}
     ```
 """
 
+import json
+import sys
 import warnings
 
 from tqdm.auto import tqdm as old_tqdm
@@ -56,6 +67,10 @@
 # By default, progress bars are enabled.
 _hf_datasets_progress_bars_disabled: bool = HF_DATASETS_DISABLE_PROGRESS_BARS or False
 
+# Progress format: "tqdm" (default), "json", or "silent"
+# Similar to huggingface/tokenizers#1921
+_hf_datasets_progress_format: str = "tqdm"
+
 
 def disable_progress_bars() -> None:
     """
@@ -101,18 +116,86 @@ def are_progress_bars_disabled() -> bool:
     return _hf_datasets_progress_bars_disabled
 
 
+def set_progress_format(format: str) -> None:
+    """
+    Set the global progress format for `datasets`.
+
+    Similar to huggingface/tokenizers#1921 progress_format option.
+
+    Args:
+        format: One of "tqdm" (default interactive bars), "json" (machine-readable JSON lines), or "silent" (no output).
+
+    Example:
+        ```py
+        from datasets.utils import set_progress_format, tqdm
+
+        # Enable JSON output for programmatic consumption
+        set_progress_format("json")
+
+        for i in tqdm(range(100), desc="Processing"):
+            do_something()
+        # Outputs: {"stage":"Processing","current":50,"total":100,"percent":50.0}
+        ```
+    """
+    if format not in ("tqdm", "json", "silent"):
+        raise ValueError(f"Invalid progress format: {format}. Must be 'tqdm', 'json', or 'silent'.")
+    global _hf_datasets_progress_format
+    _hf_datasets_progress_format = format
+
+
+def get_progress_format() -> str:
+    """
+    Get the current global progress format.
+
+    Returns:
+        Current progress format ("tqdm", "json", or "silent").
+    """
+    global _hf_datasets_progress_format
+    return _hf_datasets_progress_format
+
+
 class tqdm(old_tqdm):
     """
-    Class to override `disable` argument in case progress bars are globally disabled.
+    Class to override `disable` argument in case progress bars are globally disabled
+    and to emit JSON progress when format is "json".
 
-    Taken from https://github.com/tqdm/tqdm/issues/619#issuecomment-619639324.
+    Taken from https://github.com/tqdm/tqdm/issues/619#issuecomment-619639324
+    and enhanced with progress_format support (similar to huggingface/tokenizers#1921).
     """
 
     def __init__(self, *args, **kwargs):
         if are_progress_bars_disabled():
             kwargs["disable"] = True
+        elif get_progress_format() == "silent":
+            kwargs["disable"] = True
+        elif get_progress_format() == "json":
+            # Disable tqdm visual output, we'll emit JSON instead
+            kwargs["disable"] = True
+
         super().__init__(*args, **kwargs)
 
+        # Store description for JSON output
+        self._json_stage = kwargs.get("desc", "")
+        self._last_json_percent = -1
+
+    def update(self, n=1):
+        """Override update to emit JSON progress when format is 'json'."""
+        super().update(n)
+
+        if get_progress_format() == "json" and self.total:
+            current_percent = round((self.n / self.total) * 100, 1) if self.total > 0 else 0
+
+            # Emit JSON every 5% or at completion
+            if current_percent - self._last_json_percent >= 5.0 or self.n == self.total:
+                progress_data = {
+                    "stage": self._json_stage,
+                    "current": self.n,
+                    "total": self.total,
+                    "percent": current_percent
+                }
+                print(json.dumps(progress_data, ensure_ascii=False), file=sys.stderr, flush=True)
+                self._last_json_percent = current_percent
+
     def __delattr__(self, attr: str) -> None:
         """Fix for https://github.com/huggingface/datasets/issues/6066"""
         try: