47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import os
|
||
import torch
|
||
from fireredtts.modules.tokenizer.whisper_tokenizer import get_tokenizer
|
||
from fireredtts.modules.text_normalizer.normalize import TextNormalizer
|
||
|
||
|
||
DEFAULT_VOCAB_FILE = os.path.join(
|
||
os.path.dirname(os.path.realpath(__file__)), "../data/tokenizer.json"
|
||
)
|
||
|
||
|
||
class VoiceBpeTokenizer:
|
||
def __init__(self):
|
||
self.tokenizer = get_tokenizer(multilingual=True)
|
||
self.tn_engine = TextNormalizer()
|
||
|
||
def redtts_text_cleaner(self, text):
|
||
text = text.strip()
|
||
text, text_lang = self.tn_engine.tn(text)
|
||
# print("---text after tn:", text)
|
||
return text, text_lang
|
||
|
||
def encode(self, text, lang="auto"):
|
||
text, text_lang = self.redtts_text_cleaner(text=text)
|
||
if lang == "auto":
|
||
lang = text_lang
|
||
text = f"[{lang}]{text}"
|
||
return self.tokenizer.encode(text)
|
||
|
||
def decode(self, seq):
|
||
if isinstance(seq, torch.Tensor):
|
||
seq = seq.cpu().numpy()
|
||
text = self.tokenizer.decode(seq)
|
||
return text
|
||
|
||
def __len__(self):
|
||
return self.tokenizer.get_vocab_size()
|
||
|
||
def get_number_tokens(self):
|
||
return self.tokenizer.get_vocab_size()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
tok = VoiceBpeTokenizer()
|
||
codes = tok.encode("我、真是hello USA啊？谢谢你world！")
|
||
print([tok.decode([c]) for c in codes])
|