Add new text segmentation library for better TTS quality

NVIDIA · Oct 5, 2023 · 47b4ab5 · 47b4ab5
1 parent 399daf3
commit 47b4ab5
Show file tree

Hide file tree

Showing 2 changed files with 9 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,13 @@ 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(e)
+
+            self.word_segmenter = pkuseg.pkuseg().cut
         else:
             self.word_segmenter = lambda x: [x]
 

diff --git a/requirements/requirements_tts.txt b/requirements/requirements_tts.txt
@@ -1,6 +1,7 @@
 attrdict
 einops
 jieba
+pkuseg
 kornia
 librosa
 matplotlib