Merge pull request espnet#4320 from cadia-lvl/add-progress-bar

egs2/TEMPLATE/asr1/pyscripts/utils/convert_text_to_phn.py

@@ -1,15 +1,18 @@
 #!/usr/bin/env python3
 
-# Copyright 2021 Tomoki Hayashi
+# Copyright 2021 Tomoki Hayashi and Gunnar Thor
 #  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
 
 """Convert kaldi-style text into phonemized sentences."""
 
 import argparse
 import codecs
+import contextlib
 
 from joblib import delayed
 from joblib import Parallel
+from joblib import parallel
+from tqdm import tqdm
 
 from espnet2.text.cleaner import TextCleaner
 from espnet2.text.phoneme_tokenizer import PhonemeTokenizer
@@ -34,13 +37,40 @@ def main():
     text = {line.split()[0]: " ".join(line.split()[1:]) for line in lines}
     if cleaner is not None:
         text = {k: cleaner(v) for k, v in text.items()}
-    phns_list = Parallel(n_jobs=args.nj)(
-        [delayed(phoneme_tokenizer.text2tokens)(sentence) for sentence in text.values()]
-    )
+    with tqdm_joblib(tqdm(total=len(text.values()), desc="Phonemizing")):
+        phns_list = Parallel(n_jobs=args.nj)(
+            [
+                delayed(phoneme_tokenizer.text2tokens)(sentence)
+                for sentence in text.values()
+            ]
+        )
     with codecs.open(args.out_text, "w", encoding="utf8") as g:
         for utt_id, phns in zip(text.keys(), phns_list):
             g.write(f"{utt_id} " + " ".join(phns) + "\n")
 
 
+@contextlib.contextmanager
+def tqdm_joblib(tqdm_object):
+    """Patch joblib to report into tqdm progress bar given as argument.
+
+    Reference:
+        https://stackoverflow.com/questions/24983493
+
+    """
+
+    class TqdmBatchCompletionCallback(parallel.BatchCompletionCallBack):
+        def __call__(self, *args, **kwargs):
+            tqdm_object.update(n=self.batch_size)
+            return super().__call__(*args, **kwargs)
+
+    old_batch_callback = parallel.BatchCompletionCallBack
+    parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
+    try:
+        yield tqdm_object
+    finally:
+        parallel.BatchCompletionCallBack = old_batch_callback
+        tqdm_object.close()
+
+
 if __name__ == "__main__":
     main()