[Wav2Vec2ProcessorWithLM] improve multi processing (huggingface#15247)

patrickvonplaten · web-flow · commit 80af1048cf34 · 2022-01-21T18:30:10.000+01:00
* [Wav2Vec2ProcessorWithLM] improve multi processing

* close pool
diff --git a/src/transformers/models/wav2vec2_with_lm/processing_wav2vec2_with_lm.py b/src/transformers/models/wav2vec2_with_lm/processing_wav2vec2_with_lm.py
@@ -18,7 +18,7 @@
 import os
 from contextlib import contextmanager
 from dataclasses import dataclass
-from multiprocessing import Pool
+from multiprocessing import get_context
 from typing import TYPE_CHECKING, Iterable, List, Optional, Union
 
 import numpy as np
@@ -300,7 +300,7 @@ def batch_decode(
 
         # create multiprocessing pool and list numpy arrays
         logits_list = [array for array in logits]
-        pool = Pool(num_processes)
+        pool = get_context("fork").Pool(num_processes)
 
         # pyctcdecode
         decoded_beams = self.decoder.decode_beams_batch(
@@ -313,6 +313,9 @@ def batch_decode(
             hotword_weight=hotword_weight,
         )
 
+        # clone multi-processing pool
+        pool.close()
+
         # extract text
         batch_texts = [d[0][0] for d in decoded_beams]
 
