diff --git a/datasets/flwr_datasets/partitioner/natural_id_partitioner.py b/datasets/flwr_datasets/partitioner/natural_id_partitioner.py index 86827c748be9..5223c964a55e 100644 --- a/datasets/flwr_datasets/partitioner/natural_id_partitioner.py +++ b/datasets/flwr_datasets/partitioner/natural_id_partitioner.py @@ -18,6 +18,7 @@ from typing import Dict, Union import numpy as np +from tqdm import tqdm import datasets from flwr_datasets.common.typing import NDArrayInt @@ -99,7 +100,9 @@ def _create_partition_id_to_indices(self) -> None: unique_natural_ids, inverse = np.unique(natural_ids, return_inverse=True) - for i, natural_id in enumerate(unique_natural_ids): + for i, natural_id in tqdm( + enumerate(unique_natural_ids), desc="Generating partition_id_to_indices" + ): if none_present and natural_id == none_replacement: # Use the natural_id that is present in the dataset (not replacement) natural_id = None diff --git a/datasets/pyproject.toml b/datasets/pyproject.toml index 7dfa60138582..c16389e1529b 100644 --- a/datasets/pyproject.toml +++ b/datasets/pyproject.toml @@ -58,6 +58,7 @@ datasets = "^2.14.6" pillow = { version = ">=6.2.1", optional = true } soundfile = { version = ">=0.12.1", optional = true } librosa = { version = ">=0.10.0.post2", optional = true } +tqdm ="^4.66.1" [tool.poetry.dev-dependencies] isort = "==5.13.2"