Use tqdm from tqdm_utils

datasets/ncslgr/ncslgr.py

@@ -19,8 +19,6 @@
 import re
 from dataclasses import dataclass
 
-from tqdm import tqdm
-
 import datasets
 
 
@@ -128,7 +126,7 @@ def get_tier_values(name: str):
     def _generate_examples(self, eaf_path: str, videos_path: str):
         """Yields examples."""
 
-        for i, eaf_file in enumerate(tqdm(os.listdir(eaf_path))):
+        for i, eaf_file in enumerate(os.listdir(eaf_path)):
             eaf_file_path = os.path.join(eaf_path, eaf_file)
             videos = []
             with open(eaf_file_path, "r", encoding="utf-8") as f:

src/datasets/__init__.py

@@ -74,7 +74,6 @@
     percent,
 )
 from .utils import *
-from .utils.tqdm_utils import disable_progress_bar
 
 
 SCRIPTS_VERSION = "master" if __version__.split(".")[-1].startswith("dev") else __version__

src/datasets/arrow_dataset.py

@@ -41,7 +41,7 @@
 
 from datasets.tasks.text_classification import TextClassification
 
-from . import config
+from . import config, utils
 from .arrow_reader import ArrowReader
 from .arrow_writer import ArrowWriter, OptimizedTypedSequence
 from .features import ClassLabel, Features, Value, cast_to_python_objects
@@ -1701,7 +1701,10 @@ def format_cache_file_name(cache_file_name, rank):
             ):
                 logger.warning("Setting TOKENIZERS_PARALLELISM=false for forked processes.")
             os.environ["TOKENIZERS_PARALLELISM"] = "false"
-            with Pool(num_proc, initargs=(RLock(),), initializer=tqdm.set_lock) as pool:
+            initargs, initializer = None, None
+            if utils.is_progress_bar_enabled():
+                initargs, initializer = (RLock(),), tqdm.set_lock
+            with Pool(num_proc, initargs=initargs, initializer=initializer) as pool:
                 os.environ = prev_env
                 shards = [
                     self.shard(num_shards=num_proc, index=rank, contiguous=True, keep_in_memory=keep_in_memory)
@@ -1812,7 +1815,7 @@ def _map_single(
             logging.set_verbosity_warning()
         # Print at least one thing to fix tqdm in notebooks in multiprocessing
         # see https://github.com/tqdm/tqdm/issues/485#issuecomment-473338308
-        if rank is not None and "notebook" in tqdm.__name__:
+        if rank is not None and utils.is_progress_bar_enabled() and "notebook" in tqdm.__name__:
             print(" ", end="", flush=True)
 
         # Select the columns (arrow columns) to process
@@ -1976,7 +1979,7 @@ def init_buffer_and_writer():
                 pbar_iterable = input_dataset if not batched else range(0, len(input_dataset), batch_size)
                 pbar_unit = "ex" if not batched else "ba"
                 pbar_desc = (desc or "") + " #" + str(rank) if rank is not None else desc
-                pbar = tqdm(
+                pbar = utils.tqdm(
                     pbar_iterable,
                     disable=bool(logging.get_verbosity() == logging.NOTSET),
                     position=rank,

src/datasets/arrow_writer.py

@@ -21,9 +21,8 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pyarrow as pa
-from tqdm.auto import tqdm
 
-from . import config
+from . import config, utils
 from .features import Features, _ArrayXDExtensionType
 from .info import DatasetInfo
 from .keyhash import DuplicatedKeysError, KeyHasher
@@ -538,9 +537,9 @@ def parquet_to_arrow(sources, destination):
     stream = None if isinstance(destination, str) else destination
     disable = bool(logging.get_verbosity() == logging.NOTSET)
     with ArrowWriter(path=destination, stream=stream) as writer:
-        for source in tqdm(sources, unit="sources", disable=disable):
+        for source in utils.tqdm(sources, unit="sources", disable=disable):
             pf = pa.parquet.ParquetFile(source)
-            for i in tqdm(range(pf.num_row_groups), unit="row_groups", leave=False, disable=disable):
+            for i in utils.tqdm(range(pf.num_row_groups), unit="row_groups", leave=False, disable=disable):
                 df = pf.read_row_group(i).to_pandas()
                 for col in df.columns:
                     df[col] = df[col].apply(json.loads)

src/datasets/io/json.py

@@ -1,10 +1,10 @@
 import os
 from typing import BinaryIO, Optional, Union
 
-from .. import Dataset, Features, NamedSplit, config
+from .. import Dataset, Features, NamedSplit, config, utils
 from ..formatting import query_table
 from ..packaged_modules.json.json import Json
-from ..utils.tqdm_utils import tqdm
+from ..utils import logging
 from ..utils.typing import NestedDataStructureLike, PathLike
 from .abc import AbstractDatasetReader
 
@@ -96,7 +96,9 @@ def _write(
         written = 0
         _ = to_json_kwargs.pop("path_or_buf", None)
 
-        for offset in tqdm(range(0, len(self.dataset), batch_size)):
+        for offset in utils.tqdm(
+            range(0, len(self.dataset), batch_size), unit="ba", disable=bool(logging.get_verbosity() == logging.NOTSET)
+        ):
             batch = query_table(
                 table=self.dataset.data,
                 key=slice(offset, offset + batch_size),

src/datasets/search.py

@@ -5,8 +5,8 @@
 from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Union
 
 import numpy as np
-from tqdm.auto import tqdm
 
+from . import utils
 from .utils import logging
 
 
@@ -141,7 +141,9 @@ def add_documents(self, documents: Union[List[str], "Dataset"], column: Optional
         index_config = self.es_index_config
         self.es_client.indices.create(index=index_name, body=index_config)
         number_of_docs = len(documents)
-        progress = tqdm(unit="docs", total=number_of_docs, disable=bool(logging.get_verbosity() == logging.NOTSET))
+        progress = utils.tqdm(
+            unit="docs", total=number_of_docs, disable=bool(logging.get_verbosity() == logging.NOTSET)
+        )
         successes = 0
 
         def passage_generator():
@@ -287,7 +289,9 @@ def add_vectors(
 
         # Add vectors
         logger.info("Adding {} vectors to the faiss index".format(len(vectors)))
-        for i in tqdm(range(0, len(vectors), batch_size), disable=bool(logging.get_verbosity() == logging.NOTSET)):
+        for i in utils.tqdm(
+            range(0, len(vectors), batch_size), disable=bool(logging.get_verbosity() == logging.NOTSET)
+        ):
             vecs = vectors[i : i + batch_size] if column is None else vectors[i : i + batch_size][column]
             self.faiss_index.add(vecs)
 

src/datasets/utils/__init__.py

@@ -35,5 +35,5 @@
     zip_dict,
     zip_nested,
 )
-from .tqdm_utils import async_tqdm, tqdm
+from .tqdm_utils import async_tqdm, disable_progress_bar, is_progress_bar_enabled, tqdm
 from .version import Version

src/datasets/utils/file_utils.py

@@ -24,9 +24,8 @@
 import numpy as np
 import posixpath
 import requests
-from tqdm.auto import tqdm
 
-from .. import __version__, config
+from .. import __version__, config, utils
 from . import logging
 from .extract import ExtractManager
 from .filelock import FileLock
@@ -431,7 +430,7 @@ def http_get(url, temp_file, proxies=None, resume_size=0, headers=None, cookies=
         return
     content_length = response.headers.get("Content-Length")
     total = resume_size + int(content_length) if content_length is not None else None
-    progress = tqdm(
+    progress = utils.tqdm(
         unit="B",
         unit_scale=True,
         total=total,

src/datasets/utils/py_utils.py

@@ -33,9 +33,10 @@
 
 import dill
 import numpy as np
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
-from .logging import INFO, WARNING, get_logger, get_verbosity, set_verbosity_warning
+from .. import utils
+from . import logging
 
 
 try:  # pragma: no branch
@@ -45,7 +46,7 @@
     _typing_extensions = Literal = Final = None
 
 
-logger = get_logger(__name__)
+logger = logging.get_logger(__name__)
 
 
 # NOTE: When used on an instance method, the cache is shared across all
@@ -142,17 +143,17 @@ def _single_map_nested(args):
         return function(data_struct)
 
     # Reduce logging to keep things readable in multiprocessing with tqdm
-    if rank is not None and get_verbosity() < WARNING:
-        set_verbosity_warning()
+    if rank is not None and logging.get_verbosity() < logging.WARNING:
+        logging.set_verbosity_warning()
     # Print at least one thing to fix tqdm in notebooks in multiprocessing
     # see https://github.com/tqdm/tqdm/issues/485#issuecomment-473338308
-    if rank is not None and "notebook" in tqdm.__name__:
+    if rank is not None and utils.is_progress_bar_enabled() and "notebook" in tqdm.__name__:
         print(" ", end="", flush=True)
 
     # Loop over single examples or batches and write to buffer/file if examples are to be updated
     pbar_iterable = data_struct.items() if isinstance(data_struct, dict) else data_struct
     pbar_desc = "#" + str(rank) if rank is not None else None
-    pbar = tqdm(pbar_iterable, disable=disable_tqdm, position=rank, unit="obj", desc=pbar_desc)
+    pbar = utils.tqdm(pbar_iterable, disable=disable_tqdm, position=rank, unit="obj", desc=pbar_desc)
 
     if isinstance(data_struct, dict):
         return {k: _single_map_nested((function, v, types, None, True)) for k, v in pbar}
@@ -194,14 +195,15 @@ def map_nested(
     if not isinstance(data_struct, dict) and not isinstance(data_struct, types):
         return function(data_struct)
 
-    disable_tqdm = bool(logger.getEffectiveLevel() > INFO)
+    disable_tqdm = bool(logger.getEffectiveLevel() > logging.INFO)
     iterable = list(data_struct.values()) if isinstance(data_struct, dict) else data_struct
 
     if num_proc is None:
         num_proc = 1
     if num_proc <= 1 or len(iterable) <= num_proc:
         mapped = [
-            _single_map_nested((function, obj, types, None, True)) for obj in tqdm(iterable, disable=disable_tqdm)
+            _single_map_nested((function, obj, types, None, True))
+            for obj in utils.tqdm(iterable, disable=disable_tqdm)
         ]
     else:
         split_kwds = []  # We organize the splits ourselve (contiguous splits)
@@ -221,7 +223,10 @@ def map_nested(
                 num_proc, len(iterable), [len(i[1]) for i in split_kwds]
             )
         )
-        with Pool(num_proc, initargs=(RLock(),), initializer=tqdm.set_lock) as pool:
+        initargs, initializer = None, None
+        if utils.is_progress_bar_enabled():
+            initargs, initializer = (RLock(),), tqdm.set_lock
+        with Pool(num_proc, initargs=initargs, initializer=initializer) as pool:
             mapped = pool.map(_single_map_nested, split_kwds)
         logger.info("Finished {} processes".format(num_proc))
         mapped = [obj for proc_res in mapped for obj in proc_res]

src/datasets/utils/tqdm_utils.py

@@ -63,8 +63,13 @@ def async_tqdm(*args, **kwargs):
         return EmptyTqdm(*args, **kwargs)
 
 
+def is_progress_bar_enabled():
+    global _active
+    return bool(_active)
+
+
 def disable_progress_bar():
-    """Disabled Tqdm progress bar.
+    """Disable tqdm progress bar.
 
     Usage:
 
@@ -103,7 +108,7 @@ def _async_tqdm(*args, **kwargs):
 
 
 class _TqdmPbarAsync:
-    """Wrapper around Tqdm pbar which be shared between thread."""
+    """Wrapper around Tqdm pbar which can be shared between thread."""
 
     _tqdm_bars = []