Make get_idf_dict compatible with DDP (#140)

Tiiiger · Oct 14, 2022 · a115707 · a115707
1 parent 74f675b
commit a115707
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/bert_score/utils.py b/bert_score/utils.py
@@ -344,8 +344,11 @@ def get_idf_dict(arr, tokenizer, nthreads=4):
 
     process_partial = partial(process, tokenizer=tokenizer)
 
-    with Pool(nthreads) as p:
-        idf_count.update(chain.from_iterable(p.map(process_partial, arr)))
+    if nthreads > 0:
+        with Pool(nthreads) as p:
+            idf_count.update(chain.from_iterable(p.map(process_partial, arr)))
+    else:
+        idf_count.update(chain.from_iterable(map(process_partial, arr)))
 
     idf_dict = defaultdict(lambda: log((num_docs + 1) / (1)))
     idf_dict.update({idx: log((num_docs + 1) / (c + 1)) for (idx, c) in idf_count.items()})