diff --git a/giant/cli.py b/giant/cli.py index 2fedc4b..9cc42a5 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -9,6 +9,7 @@ from typing_extensions import Annotated import pyarrow as pa import pyarrow.parquet as pq +from tqdm import tqdm from giant import config as gconfig from giant.constants import ( @@ -343,8 +344,10 @@ def predict( writer: pq.ParquetWriter | None = None total = 0 + total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files) chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks + bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True) for path in files: for chunk in chunk_iter(path): N = len(chunk["event_id"]) @@ -455,7 +458,9 @@ def predict( writer = pq.ParquetWriter(out, table.schema) writer.write_table(table) total += N + bar.update(N) + bar.close() if writer is not None: writer.close()