[Fix] Fix #692

open-compass · Dec 24, 2024 · e57011e · e57011e
1 parent 8d14149
commit e57011e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/vlmeval/tools.py b/vlmeval/tools.py
@@ -388,6 +388,7 @@ def parse_args_eval():
     parser.add_argument('data_file', type=str)
     parser.add_argument('--judge', type=str, default=None)
     parser.add_argument('--nproc', type=int, default=4)
+    parser.add_argument('--retry', type=int, default=None)
     args = parser.parse_args()
     return args
 
@@ -458,6 +459,8 @@ def extract_dataset(file_name):
             kwargs = {'nproc': args.nproc}
             if args.judge is not None:
                 kwargs['model'] = args.judge
+            if args.retry is not None:
+                kwargs['retry'] = args.retry
             EVAL(dataset_name=dataset, data_file=data_file, **kwargs)
     else:
         logger.error('WARNING: command error!')

diff --git a/vlmeval/utils/mp_util.py b/vlmeval/utils/mp_util.py
@@ -58,8 +58,9 @@ def track_progress_rich(
                     new_finished.add(idx)
                     if keys is not None:
                         res[keys[idx]] = results[idx]
-            if len(new_finished) and save is not None:
-                dump(res, save)
+            if len(new_finished):
+                if save is not None:
+                    dump(res, save)
                 pbar.update(len(new_finished))
                 for k in new_finished:
                     unfinished.remove(k)