From e790f0b530a8eda91541100b1f5dfd64ca37d862 Mon Sep 17 00:00:00 2001
From: CodingWithTim <tianleli@berkeley.edu>
Date: Mon, 11 Nov 2024 00:37:01 +0000
Subject: [PATCH] add imap

---
 fastchat/serve/monitor/clean_chat_data.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/fastchat/serve/monitor/clean_chat_data.py b/fastchat/serve/monitor/clean_chat_data.py
index 36eb8201e..42cd17327 100644
--- a/fastchat/serve/monitor/clean_chat_data.py
+++ b/fastchat/serve/monitor/clean_chat_data.py
@@ -12,6 +12,7 @@
 from functools import partial
 from math import ceil
 from datetime import datetime, timedelta
+from tqdm import tqdm
 import time
 import multiprocessing as mp
 
@@ -139,8 +140,14 @@ def clean_chat_data(log_files, action_type, num_parallel):
     with mp.Pool(num_parallel) as pool:
         # Use partial to pass action_type to get_action_type_data
         func = partial(get_action_type_data, action_type=action_type)
-        file_data = pool.map(
-            func, log_files, chunksize=ceil(len(log_files) / len(pool._pool))
+        file_data = list(
+            tqdm(
+                pool.imap(
+                    func, log_files, chunksize=ceil(len(log_files) / len(pool._pool))
+                ),
+                total=len(log_files),
+                desc="Processing Log Files",
+            )
         )
     # filter out Nones as some files may not contain any data belong to action_type
     raw_data = []
@@ -151,8 +158,14 @@ def clean_chat_data(log_files, action_type, num_parallel):
     # Use the multiprocessing Pool
     with mp.Pool(num_parallel) as pool:
         func = partial(process_data, action_type=action_type)
-        results = pool.map(
-            func, raw_data, chunksize=ceil(len(raw_data) / len(pool._pool))
+        results = list(
+            tqdm(
+                pool.imap(
+                    func, raw_data, chunksize=ceil(len(raw_data) / len(pool._pool))
+                ),
+                total=len(raw_data),
+                desc="Processing Raw Data",
+            )
         )
 
     # Aggregate results from child processes
@@ -161,7 +174,7 @@ def clean_chat_data(log_files, action_type, num_parallel):
     ct_network_error = 0
     all_models = set()
     chats = []
-    for data in results:
+    for data in tqdm(results):
         if "ct_invalid_conv_id" in data:
             ct_invalid_conv_id += data["ct_invalid_conv_id"]
             continue