File size: 4,303 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
"""Utility helpers to handle progress bars in `datasets`.
Example:
1. Use `datasets.utils.tqdm` as you would use `tqdm.tqdm` or `tqdm.auto.tqdm`.
2. To disable progress bars, either use `disable_progress_bars()` helper or set the
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()`.
NOTE: Environment variable `HF_DATASETS_DISABLE_PROGRESS_BARS` has the priority.
Example:
```py
from datasets.utils import (
are_progress_bars_disabled,
disable_progress_bars,
enable_progress_bars,
tqdm,
)
# Disable progress bars globally
disable_progress_bars()
# Use as normal `tqdm`
for _ in tqdm(range(5)):
do_something()
# Still not showing progress bars, as `disable=False` is overwritten to `True`.
for _ in tqdm(range(5), disable=False):
do_something()
are_progress_bars_disabled() # True
# Re-enable progress bars globally
enable_progress_bars()
# Progress bar will be shown !
for _ in tqdm(range(5)):
do_something()
```
"""
import warnings
from tqdm.auto import tqdm as old_tqdm
from ..config import HF_DATASETS_DISABLE_PROGRESS_BARS
# `HF_DATASETS_DISABLE_PROGRESS_BARS` is `Optional[bool]` while `_hf_datasets_progress_bars_disabled`
# is a `bool`. If `HF_DATASETS_DISABLE_PROGRESS_BARS` is set to True or False, it has priority.
# If `HF_DATASETS_DISABLE_PROGRESS_BARS` is None, it means the user have not set the
# environment variable and is free to enable/disable progress bars programmatically.
# TL;DR: env variable has priority over code.
#
# By default, progress bars are enabled.
_hf_datasets_progress_bars_disabled: bool = HF_DATASETS_DISABLE_PROGRESS_BARS or False
def disable_progress_bars() -> None:
"""
Disable globally progress bars used in `datasets` except if `HF_DATASETS_DISABLE_PROGRESS_BAR` environment
variable has been set.
Use [`~utils.enable_progress_bars`] to re-enable them.
"""
if HF_DATASETS_DISABLE_PROGRESS_BARS is False:
warnings.warn(
"Cannot disable progress bars: environment variable `HF_DATASETS_DISABLE_PROGRESS_BAR=0` is set and has"
" priority."
)
return
global _hf_datasets_progress_bars_disabled
_hf_datasets_progress_bars_disabled = True
def enable_progress_bars() -> None:
"""
Enable globally progress bars used in `datasets` except if `HF_DATASETS_DISABLE_PROGRESS_BAR` environment
variable has been set.
Use [`~utils.disable_progress_bars`] to disable them.
"""
if HF_DATASETS_DISABLE_PROGRESS_BARS is True:
warnings.warn(
"Cannot enable progress bars: environment variable `HF_DATASETS_DISABLE_PROGRESS_BAR=1` is set and has"
" priority."
)
return
global _hf_datasets_progress_bars_disabled
_hf_datasets_progress_bars_disabled = False
def are_progress_bars_disabled() -> bool:
"""Return whether progress bars are globally disabled or not.
Progress bars used in `datasets` can be enable or disabled globally using [`~utils.enable_progress_bars`]
and [`~utils.disable_progress_bars`] or by setting `HF_DATASETS_DISABLE_PROGRESS_BAR` as environment variable.
"""
global _hf_datasets_progress_bars_disabled
return _hf_datasets_progress_bars_disabled
class tqdm(old_tqdm):
"""
Class to override `disable` argument in case progress bars are globally disabled.
Taken from https://github.com/tqdm/tqdm/issues/619#issuecomment-619639324.
"""
def __init__(self, *args, **kwargs):
if are_progress_bars_disabled():
kwargs["disable"] = True
super().__init__(*args, **kwargs)
def __delattr__(self, attr: str) -> None:
"""Fix for https://github.com/huggingface/datasets/issues/6066"""
try:
super().__delattr__(attr)
except AttributeError:
if attr != "_lock":
raise
# backward compatibility
enable_progress_bar = enable_progress_bars
disable_progress_bar = disable_progress_bars
def is_progress_bar_enabled():
return not are_progress_bars_disabled()
|