Add progress bar to serialize

app/utils.py

@@ -11,16 +11,23 @@ if TYPE_CHECKING:
 __all__ = ["serialize", "deserialize"]
 
 
-def serialize(data: Sequence[str], path: Path, max_size: int = 100000) -> None:
+def serialize(data: Sequence[str], path: Path, max_size: int = 100000, show_progress: bool = False) -> None:
     """Serialize data to a file
 
     Args:
         data: The data to serialize
         path: The path to save the serialized data
         max_size: The maximum size a chunk can be (in elements)
+        show_progress: Whether to show a progress bar
     """
     # first file is path, next chunks have ".1", ".2", etc. appended
-    for i, chunk in enumerate(tqdm([data[i : i + max_size] for i in range(0, len(data), max_size)])):
+    for i, chunk in enumerate(
+        tqdm(
+            [data[i : i + max_size] for i in range(0, len(data), max_size)],
+            unit="chunk",
+            disable=not show_progress,
+        ),
+    ):
         fd = path.with_suffix(f".{i}.pkl" if i else ".pkl")
         with fd.open("wb") as f:
             joblib.dump(chunk, f, compress=3)