Repository: huggingface/datasets
Pull request: Better TQDM output
This PR does the following:
* if `dataset_infos.json` exists for a dataset, uses `num_examples` to print the total number of examples that needs to be generated (in `builder.py`)
* fixes `tqdm` + multiprocessing in Jupyter Notebook/Colab (the issue stems from this commit in the `tqdm` repo: https://github.com/tqdm/tqdm/commit/f7722edecc3010cb35cc1c923ac4850a76336f82) 
* adds the missing `drop_last_batch` and `with_ranks` params to `DatasetDict.map` 
* correctly computes the number of iterations in `map` and the CSV/JSON loader when `batched=True` to fix `tqdm` progress bars
* removes the `bool(logging.get_verbosity() == logging.NOTSET)` (or simplifies `bool(logging.get_verbosity() == logging.NOTSET) or not utils.is_progress_bar_enabled()` to `not utils.is_progress_bar_enabled()`)  condition and uses `utils.is_progress_bar_enabled` to check if `tqdm` output is enabled (this comment from @stas00 explains why  the  `bool(logging.get_verbosity() == logging.NOTSET)` check is problematic: https://github.com/huggingface/transformers/issues/14889#issue-1087318463)

Fix #2630

Problem statement:
Progress bars are not properly rendered in Jupyter notebook
## Describe the bug
The progress bars are not Jupyter widgets; regular progress bars appear (like in a terminal).

## Steps to reproduce the bug
```python
ds.map(tokenize, num_proc=10)
```

## Expected results
Jupyter widgets displaying the progress bars.

## Actual results
Simple plane progress bars.

cc: Reported by @thomwolf

Hints from issue discussion:
To add my experience when trying to debug this issue:

Seems like previously the workaround given [here](https://github.com/tqdm/tqdm/issues/485#issuecomment-473338308) worked around this issue. But with the latest version of jupyter/tqdm I still get terminal warnings that IPython tried to send a message from a forked process.
Hi @mludv, thanks for the hint!!! :) 

We will definitely take it into account to try to fix this issue... It seems somehow related to `multiprocessing` and `tqdm`...

Commit message: Add more precise descriptions to download

Patch to review:
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index 89a68fb57bd..401d2b9115f 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -2096,7 +2096,7 @@ def decorated(item, *args, **kwargs):
                 f"num_proc must be <= {len(self)}. Reducing num_proc to {num_proc} for dataset of size {len(self)}."
             )
 
-        disable_tqdm = bool(logging.get_verbosity() == logging.NOTSET) or not utils.is_progress_bar_enabled()
+        disable_tqdm = not utils.is_progress_bar_enabled()
 
         if num_proc is None or num_proc == 1:
             return self._map_single(
@@ -2300,7 +2300,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 not disable_tqdm and "notebook" in tqdm.__name__:
+        if rank is not None and not disable_tqdm and any("notebook" in tqdm_cls.__name__ for tqdm_cls in tqdm.__mro__):
             print(" ", end="", flush=True)
 
         if fn_kwargs is None:
@@ -2443,13 +2443,20 @@ def init_buffer_and_writer():
                     input_dataset = self
 
                 # Loop over single examples or batches and write to buffer/file if examples are to be updated
-                pbar_iterable = (
-                    input_dataset._iter(decoded=False) if not batched else range(0, len(input_dataset), batch_size)
-                )
+                if not batched:
+                    pbar_iterable = input_dataset._iter(decoded=False)
+                    pbar_total = len(input_dataset)
+                else:
+                    num_rows = (
+                        len(input_dataset) if not drop_last_batch else len(input_dataset) // batch_size * batch_size
+                    )
+                    pbar_iterable = range(0, num_rows, batch_size)
+                    pbar_total = (num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size
                 pbar_unit = "ex" if not batched else "ba"
                 pbar_desc = (desc or "") + " #" + str(rank) if rank is not None else desc
                 pbar = utils.tqdm(
                     pbar_iterable,
+                    total=pbar_total,
                     disable=disable_tqdm,
                     position=rank,
                     unit=pbar_unit,
@@ -2468,8 +2475,6 @@ def init_buffer_and_writer():
                                 writer.write(example)
                 else:
                     for i in pbar:
-                        if drop_last_batch and i + batch_size > input_dataset.num_rows:
-                            continue
                         batch = input_dataset._getitem(
                             slice(i, i + batch_size),
                             decoded=False,
@@ -3539,7 +3544,7 @@ def delete_file(file):
                 file_shards_to_delete,
                 desc="Deleting unused files from dataset repository",
                 total=len(file_shards_to_delete),
-                disable=bool(logging.get_verbosity() == logging.NOTSET) or not utils.is_progress_bar_enabled(),
+                disable=not utils.is_progress_bar_enabled(),
             ):
                 delete_file(file)
 
@@ -3548,7 +3553,7 @@ def delete_file(file):
             enumerate(shards),
             desc="Pushing dataset shards to the dataset hub",
             total=num_shards,
-            disable=bool(logging.get_verbosity() == logging.NOTSET),
+            disable=not utils.is_progress_bar_enabled(),
         ):
             buffer = BytesIO()
             shard.to_parquet(buffer)
diff --git a/src/datasets/arrow_writer.py b/src/datasets/arrow_writer.py
index 91dd4f00662..e996c0344f7 100644
--- a/src/datasets/arrow_writer.py
+++ b/src/datasets/arrow_writer.py
@@ -640,7 +640,7 @@ def finalize(self, metrics_query_result: dict):
 def parquet_to_arrow(sources, destination):
     """Convert parquet files to arrow file. Inputs can be str paths or file-like objects"""
     stream = None if isinstance(destination, str) else destination
-    disable = bool(logging.get_verbosity() == logging.NOTSET)
+    disable = not utils.is_progress_bar_enabled()
     with ArrowWriter(path=destination, stream=stream) as writer:
         for source in utils.tqdm(sources, unit="sources", disable=disable):
             pf = pa.parquet.ParquetFile(source)
diff --git a/src/datasets/builder.py b/src/datasets/builder.py
index 668f197e936..d14a8351bf6 100644
--- a/src/datasets/builder.py
+++ b/src/datasets/builder.py
@@ -675,7 +675,7 @@ def _download_and_prepare(self, dl_manager, verify_infos, **prepare_split_kwargs
                     "._split_generator()."
                 )
 
-            logger.info(f"Generating split {split_generator.split_info.name}")
+            logger.info(f"Generating {split_generator.split_info.name} split")
             split_dict.add(split_generator.split_info)
 
             try:
@@ -769,7 +769,7 @@ def as_dataset(
             ),
             split,
             map_tuple=True,
-            disable_tqdm=False,
+            disable_tqdm=not utils.is_progress_bar_enabled(),
         )
         if isinstance(datasets, dict):
             datasets = DatasetDict(datasets)
@@ -1055,7 +1055,10 @@ def _generate_examples(self, **kwargs):
         raise NotImplementedError()
 
     def _prepare_split(self, split_generator):
-        split_info = split_generator.split_info
+        if self.info.splits is not None:
+            split_info = self.info.splits[split_generator.name]
+        else:
+            split_info = split_generator.split_info
 
         fname = f"{self.name}-{split_generator.name}.arrow"
         fpath = os.path.join(self._cache_dir, fname)
@@ -1075,7 +1078,8 @@ def _prepare_split(self, split_generator):
                     unit=" examples",
                     total=split_info.num_examples,
                     leave=False,
-                    disable=bool(logging.get_verbosity() == logging.NOTSET),
+                    disable=not utils.is_progress_bar_enabled(),
+                    desc=f"Generating {split_info.name} split",
                 ):
                     example = self.info.features.encode_example(record)
                     writer.write(example, key)
@@ -1131,7 +1135,7 @@ def _prepare_split(self, split_generator):
         generator = self._generate_tables(**split_generator.gen_kwargs)
         with ArrowWriter(features=self.info.features, path=fpath) as writer:
             for key, table in utils.tqdm(
-                generator, unit=" tables", leave=False, disable=True  # bool(logging.get_verbosity() == logging.NOTSET)
+                generator, unit=" tables", leave=False, disable=True  # not utils.is_progress_bar_enabled()
             ):
                 writer.write_table(table)
             num_examples, num_bytes = writer.finalize()
diff --git a/src/datasets/data_files.py b/src/datasets/data_files.py
index e639e8fd367..b87c349d9c8 100644
--- a/src/datasets/data_files.py
+++ b/src/datasets/data_files.py
@@ -13,7 +13,7 @@
 from .utils import logging
 from .utils.file_utils import hf_hub_url, is_remote_url, request_etag
 from .utils.py_utils import string_to_dict
-from .utils.tqdm_utils import tqdm
+from .utils.tqdm_utils import is_progress_bar_enabled, tqdm
 
 
 DEFAULT_SPLIT = str(Split.TRAIN)
@@ -497,7 +497,7 @@ def _get_origin_metadata_locally_or_by_urls(
         max_workers=max_workers,
         tqdm_class=tqdm,
         desc="Resolving data files",
-        disable=len(data_files) <= 16 or logging.get_verbosity() == logging.NOTSET,
+        disable=len(data_files) <= 16 or not is_progress_bar_enabled(),
     )
 
 
diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py
index 6f196eabdc4..9b5863c29d5 100644
--- a/src/datasets/dataset_dict.py
+++ b/src/datasets/dataset_dict.py
@@ -435,11 +435,13 @@ def with_transform(
 
     def map(
         self,
-        function,
+        function: Optional[Callable] = None,
         with_indices: bool = False,
+        with_rank: bool = False,
         input_columns: Optional[Union[str, List[str]]] = None,
         batched: bool = False,
         batch_size: Optional[int] = 1000,
+        drop_last_batch: bool = False,
         remove_columns: Optional[Union[str, List[str]]] = None,
         keep_in_memory: bool = False,
         load_from_cache_file: bool = True,
@@ -462,11 +464,15 @@ def map(
                 - `function(batch: Dict[List]) -> Union[Dict, Any]` if `batched=True` and `with_indices=False`
                 - `function(batch: Dict[List], indices: List[int]) -> Union[Dict, Any]` if `batched=True` and `with_indices=True`
             with_indices (`bool`, defaults to `False`): Provide example indices to `function`. Note that in this case the signature of `function` should be `def function(example, idx): ...`.
+            with_rank (:obj:`bool`, default `False`): Provide process rank to `function`. Note that in this case the
+                signature of `function` should be `def function(example[, idx], rank): ...`.
             input_columns (`Optional[Union[str, List[str]]]`, defaults to `None`): The columns to be passed into `function` as
                 positional arguments. If `None`, a dict mapping to all formatted columns is passed as one argument.
             batched (`bool`, defaults to `False`): Provide batch of examples to `function`
             batch_size (`Optional[int]`, defaults to `1000`): Number of examples per batch provided to `function` if `batched=True`
                 `batch_size <= 0` or `batch_size == None`: Provide the full dataset as a single batch to `function`
+            drop_last_batch (:obj:`bool`, default `False`): Whether a last batch smaller than the batch_size should be
+                dropped instead of being processed by the function.
             remove_columns (`Optional[Union[str, List[str]]]`, defaults to `None`): Remove a selection of columns while doing the mapping.
                 Columns will be removed before updating the examples with the output of `function`, i.e. if `function` is adding
                 columns with names in `remove_columns`, these columns will be kept.
@@ -495,9 +501,11 @@ def map(
                 k: dataset.map(
                     function=function,
                     with_indices=with_indices,
+                    with_rank=with_rank,
                     input_columns=input_columns,
                     batched=batched,
                     batch_size=batch_size,
+                    drop_last_batch=drop_last_batch,
                     remove_columns=remove_columns,
                     keep_in_memory=keep_in_memory,
                     load_from_cache_file=load_from_cache_file,
diff --git a/src/datasets/io/csv.py b/src/datasets/io/csv.py
index 6b6fd604c6f..a9d8824d75d 100644
--- a/src/datasets/io/csv.py
+++ b/src/datasets/io/csv.py
@@ -5,7 +5,6 @@
 from .. import Dataset, Features, NamedSplit, config, utils
 from ..formatting import query_table
 from ..packaged_modules.csv.csv import Csv
-from ..utils import logging
 from ..utils.typing import NestedDataStructureLike, PathLike
 from .abc import AbstractDatasetReader
 
@@ -108,22 +107,23 @@ def _write(self, file_obj: BinaryIO, header: bool = True, **to_csv_kwargs) -> in
             for offset in utils.tqdm(
                 range(0, len(self.dataset), self.batch_size),
                 unit="ba",
-                disable=bool(logging.get_verbosity() == logging.NOTSET),
+                disable=not utils.is_progress_bar_enabled(),
                 desc="Creating CSV from Arrow format",
             ):
                 csv_str = self._batch_csv((offset, header, to_csv_kwargs))
                 written += file_obj.write(csv_str)
 
         else:
+            num_rows, bsz = len(self.dataset), self.batch_size
             with multiprocessing.Pool(self.num_proc) as pool:
                 for csv_str in utils.tqdm(
                     pool.imap(
                         self._batch_csv,
-                        [(offset, header, to_csv_kwargs) for offset in range(0, len(self.dataset), self.batch_size)],
+                        [(offset, header, to_csv_kwargs) for offset in range(0, num_rows, bsz)],
                     ),
-                    total=(len(self.dataset) // self.batch_size) + 1,
+                    total=(num_rows // bsz) + 1 if num_rows % bsz else num_rows // bsz,
                     unit="ba",
-                    disable=bool(logging.get_verbosity() == logging.NOTSET),
+                    disable=not utils.is_progress_bar_enabled(),
                     desc="Creating CSV from Arrow format",
                 ):
                     written += file_obj.write(csv_str)
diff --git a/src/datasets/io/json.py b/src/datasets/io/json.py
index 890be857f79..c3b042d1b10 100644
--- a/src/datasets/io/json.py
+++ b/src/datasets/io/json.py
@@ -5,7 +5,6 @@
 from .. import Dataset, Features, NamedSplit, config, utils
 from ..formatting import query_table
 from ..packaged_modules.json.json import Json
-from ..utils import logging
 from ..utils.typing import NestedDataStructureLike, PathLike
 from .abc import AbstractDatasetReader
 from .handler import IOHandler
@@ -124,24 +123,22 @@ def _write(
             for offset in utils.tqdm(
                 range(0, len(self.dataset), self.batch_size),
                 unit="ba",
-                disable=bool(logging.get_verbosity() == logging.NOTSET),
+                disable=not utils.is_progress_bar_enabled(),
                 desc="Creating json from Arrow format",
             ):
                 json_str = self._batch_json((offset, orient, lines, to_json_kwargs))
                 written += file_obj.write(json_str)
         else:
+            num_rows, bsz = len(self.dataset), self.batch_size
             with multiprocessing.Pool(self.num_proc) as pool:
                 for json_str in utils.tqdm(
                     pool.imap(
                         self._batch_json,
-                        [
-                            (offset, orient, lines, to_json_kwargs)
-                            for offset in range(0, len(self.dataset), self.batch_size)
-                        ],
+                        [(offset, orient, lines, to_json_kwargs) for offset in range(0, num_rows, bsz)],
                     ),
-                    total=(len(self.dataset) // self.batch_size) + 1,
+                    total=(num_rows // bsz) + 1 if num_rows % bsz else num_rows // bsz,
                     unit="ba",
-                    disable=bool(logging.get_verbosity() == logging.NOTSET),
+                    disable=not utils.is_progress_bar_enabled(),
                     desc="Creating json from Arrow format",
                 ):
                     written += file_obj.write(json_str)
diff --git a/src/datasets/load.py b/src/datasets/load.py
index 59253359446..82792b38df9 100644
--- a/src/datasets/load.py
+++ b/src/datasets/load.py
@@ -259,6 +259,9 @@ def _download_additional_modules(
     """
     local_imports = []
     library_imports = []
+    download_config = download_config.copy()
+    if download_config.download_desc is None:
+        download_config.download_desc = "Downloading extra module"
     for import_type, import_name, import_path, sub_directory in imports:
         if import_type == "library":
             library_imports.append((import_name, import_path))  # Import from a library
@@ -494,7 +497,7 @@ def __init__(
     ):
         self.name = name
         self.revision = revision
-        self.download_config = download_config
+        self.download_config = download_config or DownloadConfig()
         self.download_mode = download_mode
         self.dynamic_modules_path = dynamic_modules_path
         assert self.name.count("/") == 0
@@ -502,11 +505,17 @@ def __init__(
 
     def download_loading_script(self, revision: Optional[str]) -> str:
         file_path = hf_github_url(path=self.name, name=self.name + ".py", revision=revision)
+        download_config = self.download_config.copy()
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading builder script"
         return cached_path(file_path, download_config=self.download_config)
 
     def download_dataset_infos_file(self, revision: Optional[str]) -> str:
         dataset_infos = hf_github_url(path=self.name, name=config.DATASETDICT_INFOS_FILENAME, revision=revision)
         # Download the dataset infos file if available
+        download_config = self.download_config.copy()
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading metadata"
         try:
             return cached_path(
                 dataset_infos,
@@ -569,7 +578,7 @@ def __init__(
     ):
         self.name = name
         self.revision = revision
-        self.download_config = download_config
+        self.download_config = download_config or DownloadConfig()
         self.download_mode = download_mode
         self.dynamic_modules_path = dynamic_modules_path
         assert self.name.count("/") == 0
@@ -577,6 +586,9 @@ def __init__(
 
     def download_loading_script(self, revision: Optional[str]) -> str:
         file_path = hf_github_url(path=self.name, name=self.name + ".py", revision=revision, dataset=False)
+        download_config = self.download_config.copy()
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading builder script"
         return cached_path(file_path, download_config=self.download_config)
 
     def get_module(self) -> MetricModule:
@@ -630,7 +642,7 @@ def __init__(
     ):
         self.path = path
         self.name = Path(path).stem
-        self.download_config = download_config
+        self.download_config = download_config or DownloadConfig()
         self.download_mode = download_mode
         self.dynamic_modules_path = dynamic_modules_path
 
@@ -671,7 +683,7 @@ def __init__(
     ):
         self.path = path
         self.name = Path(path).stem
-        self.download_config = download_config
+        self.download_config = download_config or DownloadConfig()
         self.download_mode = download_mode
         self.dynamic_modules_path = dynamic_modules_path
 
@@ -836,6 +848,9 @@ def get_module(self) -> DatasetModule:
             "name": self.name.replace("/", "--"),
             "base_path": hf_hub_url(self.name, "", revision=self.revision),
         }
+        download_config = self.download_config.copy()
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading metadata"
         try:
             dataset_infos_path = cached_path(
                 hf_hub_url(self.name, config.DATASETDICT_INFOS_FILENAME, revision=self.revision),
@@ -863,7 +878,7 @@ def __init__(
     ):
         self.name = name
         self.revision = revision
-        self.download_config = download_config
+        self.download_config = download_config or DownloadConfig()
         self.download_mode = download_mode
         self.dynamic_modules_path = dynamic_modules_path
         assert self.name.count("/") == 1
@@ -871,11 +886,17 @@ def __init__(
 
     def download_loading_script(self) -> str:
         file_path = hf_hub_url(path=self.name, name=self.name.split("/")[1] + ".py", revision=self.revision)
+        download_config = self.download_config.copy()
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading builder script"
         return cached_path(file_path, download_config=self.download_config)
 
     def download_dataset_infos_file(self) -> str:
         dataset_infos = hf_hub_url(path=self.name, name=config.DATASETDICT_INFOS_FILENAME, revision=self.revision)
         # Download the dataset infos file if available
+        download_config = self.download_config.copy()
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading metadata"
         try:
             return cached_path(
                 dataset_infos,
diff --git a/src/datasets/search.py b/src/datasets/search.py
index 94e3045c32f..3a16e39b6ff 100644
--- a/src/datasets/search.py
+++ b/src/datasets/search.py
@@ -150,9 +150,7 @@ 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 = utils.tqdm(
-            unit="docs", total=number_of_docs, disable=bool(logging.get_verbosity() == logging.NOTSET)
-        )
+        progress = utils.tqdm(unit="docs", total=number_of_docs, disable=not utils.is_progress_bar_enabled())
         successes = 0
 
         def passage_generator():
@@ -297,9 +295,7 @@ def add_vectors(
 
         # Add vectors
         logger.info(f"Adding {len(vectors)} vectors to the faiss index")
-        for i in utils.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=not utils.is_progress_bar_enabled()):
             vecs = vectors[i : i + batch_size] if column is None else vectors[i : i + batch_size][column]
             self.faiss_index.add(vecs)
 
diff --git a/src/datasets/utils/download_manager.py b/src/datasets/utils/download_manager.py
index d0bc1ae7a09..7fb69398a28 100644
--- a/src/datasets/utils/download_manager.py
+++ b/src/datasets/utils/download_manager.py
@@ -23,7 +23,7 @@
 from functools import partial
 from typing import Dict, Optional, Union
 
-from .. import config
+from .. import config, utils
 from .file_utils import (
     DownloadConfig,
     cached_path,
@@ -122,7 +122,9 @@ def upload(local_file_path):
             return remote_file_path
 
         uploaded_path_or_paths = map_nested(
-            lambda local_file_path: upload(local_file_path), downloaded_path_or_paths, disable_tqdm=False
+            lambda local_file_path: upload(local_file_path),
+            downloaded_path_or_paths,
+            disable_tqdm=not utils.is_progress_bar_enabled(),
         )
         return uploaded_path_or_paths
 
@@ -152,7 +154,9 @@ def download_custom(self, url_or_urls, custom_download):
         def url_to_downloaded_path(url):
             return os.path.join(cache_dir, hash_url_to_filename(url))
 
-        downloaded_path_or_paths = map_nested(url_to_downloaded_path, url_or_urls, disable_tqdm=False)
+        downloaded_path_or_paths = map_nested(
+            url_to_downloaded_path, url_or_urls, disable_tqdm=not utils.is_progress_bar_enabled()
+        )
         url_or_urls = NestedDataStructure(url_or_urls)
         downloaded_path_or_paths = NestedDataStructure(downloaded_path_or_paths)
         for url, path in zip(url_or_urls.flatten(), downloaded_path_or_paths.flatten()):
@@ -193,7 +197,12 @@ def download(self, url_or_urls):
 
         start_time = datetime.now()
         downloaded_path_or_paths = map_nested(
-            download_func, url_or_urls, map_tuple=True, num_proc=download_config.num_proc, disable_tqdm=False
+            download_func,
+            url_or_urls,
+            map_tuple=True,
+            num_proc=download_config.num_proc,
+            disable_tqdm=not utils.is_progress_bar_enabled(),
+            desc="Downloading data files",
         )
         duration = datetime.now() - start_time
         logger.info(f"Downloading took {duration.total_seconds() // 60} min")
@@ -280,8 +289,15 @@ def extract(self, path_or_paths, num_proc=None):
         """
         download_config = self.download_config.copy()
         download_config.extract_compressed_file = True
+        # Extract downloads the file first if it is not already downloaded
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading data"
         extracted_paths = map_nested(
-            partial(cached_path, download_config=download_config), path_or_paths, num_proc=num_proc, disable_tqdm=False
+            partial(cached_path, download_config=download_config),
+            path_or_paths,
+            num_proc=num_proc,
+            disable_tqdm=not utils.is_progress_bar_enabled(),
+            desc="Extracting data files",
         )
         path_or_paths = NestedDataStructure(path_or_paths)
         extracted_paths = NestedDataStructure(extracted_paths)
diff --git a/src/datasets/utils/file_utils.py b/src/datasets/utils/file_utils.py
index 1b14f7125b3..c829fdb917f 100644
--- a/src/datasets/utils/file_utils.py
+++ b/src/datasets/utils/file_utils.py
@@ -241,6 +241,7 @@ class DownloadConfig:
             for remote files on the Datasets Hub. If True, will get token from ~/.huggingface.
         ignore_url_params (:obj:`bool`, default ``False``): Whether to strip all query parameters and #fragments from
             the download URL before using it for caching the file.
+        download_description (:obj:`str`, optional): A description to be displayed alongside with the progress bar while downloading the files.
     """
 
     cache_dir: Optional[Union[str, Path]] = None
@@ -257,6 +258,7 @@ class DownloadConfig:
     max_retries: int = 1
     use_auth_token: Optional[Union[str, bool]] = None
     ignore_url_params: bool = False
+    download_desc: Optional[str] = None
 
     def copy(self) -> "DownloadConfig":
         return self.__class__(**{k: copy.deepcopy(v) for k, v in self.__dict__.items()})
@@ -307,6 +309,7 @@ def cached_path(
             max_retries=download_config.max_retries,
             use_auth_token=download_config.use_auth_token,
             ignore_url_params=download_config.ignore_url_params,
+            download_desc=download_config.download_desc,
         )
     elif os.path.exists(url_or_filename):
         # File, and it exists.
@@ -434,7 +437,9 @@ def ftp_get(url, temp_file, timeout=10.0):
         raise ConnectionError(e) from None
 
 
-def http_get(url, temp_file, proxies=None, resume_size=0, headers=None, cookies=None, timeout=100.0, max_retries=0):
+def http_get(
+    url, temp_file, proxies=None, resume_size=0, headers=None, cookies=None, timeout=100.0, max_retries=0, desc=None
+):
     headers = copy.deepcopy(headers) or {}
     headers["user-agent"] = get_datasets_user_agent(user_agent=headers.get("user-agent"))
     if resume_size > 0:
@@ -453,19 +458,17 @@ 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 = utils.tqdm(
+    with utils.tqdm(
         unit="B",
         unit_scale=True,
         total=total,
         initial=resume_size,
-        desc="Downloading",
-        disable=bool(logging.get_verbosity() == logging.NOTSET),
-    )
-    for chunk in response.iter_content(chunk_size=1024):
-        if chunk:  # filter out keep-alive new chunks
+        desc=desc or "Downloading",
+        disable=not utils.is_progress_bar_enabled(),
+    ) as progress:
+        for chunk in response.iter_content(chunk_size=1024):
             progress.update(len(chunk))
             temp_file.write(chunk)
-    progress.close()
 
 
 def http_head(
@@ -507,6 +510,7 @@ def get_from_cache(
     max_retries=0,
     use_auth_token=None,
     ignore_url_params=False,
+    download_desc=None,
 ) -> str:
     """
     Given a URL, look for the corresponding file in the local cache.
@@ -664,6 +668,7 @@ def _resumable_file_manager():
                     headers=headers,
                     cookies=cookies,
                     max_retries=max_retries,
+                    desc=download_desc,
                 )
 
         logger.info(f"storing {url} in cache at {cache_path}")
diff --git a/src/datasets/utils/py_utils.py b/src/datasets/utils/py_utils.py
index 6a4b5820591..b6b35c5ca30 100644
--- a/src/datasets/utils/py_utils.py
+++ b/src/datasets/utils/py_utils.py
@@ -189,7 +189,7 @@ def __get__(self, obj, objtype=None):
 
 def _single_map_nested(args):
     """Apply a function recursively to each element of a nested data struct."""
-    function, data_struct, types, rank, disable_tqdm = args
+    function, data_struct, types, rank, disable_tqdm, desc = args
 
     # Singleton first to spare some computation
     if not isinstance(data_struct, dict) and not isinstance(data_struct, types):
@@ -200,18 +200,18 @@ def _single_map_nested(args):
         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 not disable_tqdm and "notebook" in tqdm.__name__:
+    if rank is not None and not disable_tqdm and any("notebook" in tqdm_cls.__name__ for tqdm_cls in tqdm.__mro__):
         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_desc = (desc or "") + " #" + str(rank) if rank is not None else 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}
+        return {k: _single_map_nested((function, v, types, None, True, None)) for k, v in pbar}
     else:
-        mapped = [_single_map_nested((function, v, types, None, True)) for v in pbar]
+        mapped = [_single_map_nested((function, v, types, None, True, None)) for v in pbar]
         if isinstance(data_struct, list):
             return mapped
         elif isinstance(data_struct, tuple):
@@ -230,6 +230,7 @@ def map_nested(
     num_proc: Optional[int] = None,
     types=None,
     disable_tqdm: bool = True,
+    desc: Optional[str] = None,
 ):
     """Apply a function recursively to each element of a nested data struct.
     If num_proc > 1 and the length of data_struct is longer than num_proc: use multi-processing
@@ -249,17 +250,15 @@ def map_nested(
     if not isinstance(data_struct, dict) and not isinstance(data_struct, types):
         return function(data_struct)
 
-    disable_tqdm = (
-        disable_tqdm or bool(logging.get_verbosity() == logging.NOTSET) or not utils.is_progress_bar_enabled()
-    )
+    disable_tqdm = disable_tqdm or not utils.is_progress_bar_enabled()
     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 utils.tqdm(iterable, disable=disable_tqdm)
+            _single_map_nested((function, obj, types, None, True, None))
+            for obj in utils.tqdm(iterable, disable=disable_tqdm, desc=desc)
         ]
     else:
         split_kwds = []  # We organize the splits ourselve (contiguous splits)
@@ -268,7 +267,7 @@ def map_nested(
             mod = len(iterable) % num_proc
             start = div * index + min(index, mod)
             end = start + div + (1 if index < mod else 0)
-            split_kwds.append((function, iterable[start:end], types, index, disable_tqdm))
+            split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc))
 
         if len(iterable) != sum(len(i[1]) for i in split_kwds):
             raise ValueError(


Code review comments:
diff_hunk: @@ -200,18 +200,18 @@ def _single_map_nested(args):
         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 not disable_tqdm and "notebook" in tqdm.__name__:
+    if rank is not None and not disable_tqdm and any("notebook" in tqdm_cls.__name__ for tqdm_cls in tqdm.__mro__):
         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_desc = (desc or "") + " #" + str(rank) if rank is not None else desc
path: src/datasets/utils/py_utils.py
text: @user1:
this adds a space before the # if desc is None, but it's fine I think

@author:
This line comes from `Dataset._map_single` and I just copy-pasted it here, so I'll also remove the space there.

diff_hunk: @@ -108,22 +107,23 @@ def _write(self, file_obj: BinaryIO, header: bool = True, **to_csv_kwargs) -> in
             for offset in utils.tqdm(
                 range(0, len(self.dataset), self.batch_size),
                 unit="ba",
-                disable=bool(logging.get_verbosity() == logging.NOTSET),
+                disable=not utils.is_progress_bar_enabled(),
                 desc="Creating CSV from Arrow format",
             ):
                 csv_str = self._batch_csv((offset, header, to_csv_kwargs))
                 written += file_obj.write(csv_str)
 
         else:
+            num_rows, bsz = len(self.dataset), self.batch_size
path: src/datasets/io/csv.py
text: @user1:
(nit) I think we use `batch_size` instead of `bsz` in the rest of the lib