From 9a61ba79d2af1ca35dfd6be681ebf4bcf3235e0f Mon Sep 17 00:00:00 2001
From: Adam Narozniak <adam@flower.dev>
Date: Mon, 29 Apr 2024 13:49:09 +0200
Subject: [PATCH] Add tqdm

---
 datasets/flwr_datasets/partitioner/natural_id_partitioner.py | 5 ++++-
 datasets/pyproject.toml                                      | 1 +
 2 files changed, 5 insertions(+), 1 deletion(-)

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"