Add tqdm

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

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"