Add new text segmentation library for better TTS quality (NVIDIA#7645)

* Add new text segmentation library for better TTS quality * Update zh_cn_pinyin.py added detailed instruction on how to install pkuseg. Signed-off-by: Xuesong Yang <[email protected]> * Update requirements_tts.txt remove pkuseg as the default dependency of NeMo TTS, and instead, direct users to manually install pkuseg if they really need. Signed-off-by: Xuesong Yang <[email protected]> --------- Signed-off-by: Xuesong Yang <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Xuesong Yang <[email protected]> Signed-off-by: Piotr Żelasko <[email protected]>
pzelasko · Jan 3, 2024 · 31ad8fc · 31ad8fc
1 parent d4c06c0
commit 31ad8fc
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/nemo/collections/tts/g2p/models/zh_cn_pinyin.py b/nemo/collections/tts/g2p/models/zh_cn_pinyin.py
@@ -67,6 +67,7 @@ def __init__(
         assert word_segmenter in [
             None,
             'jieba',
+            'pkuseg',
         ], f"{word_segmenter} is not supported now. Please choose correct word_segmenter."
 
         if phoneme_prefix is None:
@@ -118,6 +119,15 @@ def __init__(
 
             # Cut sentences into words to improve polyphone disambiguation
             self.word_segmenter = jieba.cut
+        elif word_segmenter == "pkuseg":
+            try:
+                import pkuseg
+            except ImportError as e:
+                logging.error(
+                    "`pkuseg` is not the default word segmenter for Chinese in NeMo. Please install it manually by running `pip install pkuseg`, or choose `jieba` segmenter instead."
+                )
+
+            self.word_segmenter = pkuseg.pkuseg().cut
         else:
             self.word_segmenter = lambda x: [x]