preprocessor.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806
  1. import re
  2. from abc import ABC
  3. from abc import abstractmethod
  4. from pathlib import Path
  5. from typing import Collection
  6. from typing import Dict
  7. from typing import Iterable
  8. from typing import List
  9. from typing import Union
  10. import numpy as np
  11. import scipy.signal
  12. import soundfile
  13. from typeguard import check_argument_types
  14. from typeguard import check_return_type
  15. from funasr.text.build_tokenizer import build_tokenizer
  16. from funasr.text.cleaner import TextCleaner
  17. from funasr.text.token_id_converter import TokenIDConverter
  18. class AbsPreprocessor(ABC):
  19. def __init__(self, train: bool):
  20. self.train = train
  21. @abstractmethod
  22. def __call__(
  23. self, uid: str, data: Dict[str, Union[str, np.ndarray]]
  24. ) -> Dict[str, np.ndarray]:
  25. raise NotImplementedError
  26. def forward_segment(text, dic):
  27. word_list = []
  28. i = 0
  29. while i < len(text):
  30. longest_word = text[i]
  31. for j in range(i + 1, len(text) + 1):
  32. word = text[i:j]
  33. if word in dic:
  34. if len(word) > len(longest_word):
  35. longest_word = word
  36. word_list.append(longest_word)
  37. i += len(longest_word)
  38. return word_list
  39. def seg_tokenize(txt, seg_dict):
  40. out_txt = ""
  41. pattern = re.compile(r"([\u4E00-\u9FA5A-Za-z0-9])")
  42. for word in txt:
  43. if pattern.match(word):
  44. if word in seg_dict:
  45. out_txt += seg_dict[word] + " "
  46. else:
  47. out_txt += "<unk>" + " "
  48. else:
  49. continue
  50. return out_txt.strip().split()
  51. def seg_tokenize_wo_pattern(txt, seg_dict):
  52. out_txt = ""
  53. for word in txt:
  54. if word in seg_dict:
  55. out_txt += seg_dict[word] + " "
  56. else:
  57. out_txt += "<unk>" + " "
  58. return out_txt.strip().split()
  59. def framing(
  60. x,
  61. frame_length: int = 512,
  62. frame_shift: int = 256,
  63. centered: bool = True,
  64. padded: bool = True,
  65. ):
  66. if x.size == 0:
  67. raise ValueError("Input array size is zero")
  68. if frame_length < 1:
  69. raise ValueError("frame_length must be a positive integer")
  70. if frame_length > x.shape[-1]:
  71. raise ValueError("frame_length is greater than input length")
  72. if 0 >= frame_shift:
  73. raise ValueError("frame_shift must be greater than 0")
  74. if centered:
  75. pad_shape = [(0, 0) for _ in range(x.ndim - 1)] + [
  76. (frame_length // 2, frame_length // 2)
  77. ]
  78. x = np.pad(x, pad_shape, mode="constant", constant_values=0)
  79. if padded:
  80. # Pad to integer number of windowed segments
  81. # I.e make x.shape[-1] = frame_length + (nseg-1)*nstep,
  82. # with integer nseg
  83. nadd = (-(x.shape[-1] - frame_length) % frame_shift) % frame_length
  84. pad_shape = [(0, 0) for _ in range(x.ndim - 1)] + [(0, nadd)]
  85. x = np.pad(x, pad_shape, mode="constant", constant_values=0)
  86. # Created strided array of data segments
  87. if frame_length == 1 and frame_length == frame_shift:
  88. result = x[..., None]
  89. else:
  90. shape = x.shape[:-1] + (
  91. (x.shape[-1] - frame_length) // frame_shift + 1,
  92. frame_length,
  93. )
  94. strides = x.strides[:-1] + (frame_shift * x.strides[-1], x.strides[-1])
  95. result = np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides)
  96. return result
  97. def detect_non_silence(
  98. x: np.ndarray,
  99. threshold: float = 0.01,
  100. frame_length: int = 1024,
  101. frame_shift: int = 512,
  102. window: str = "boxcar",
  103. ) -> np.ndarray:
  104. """Power based voice activity detection.
  105. Args:
  106. x: (Channel, Time)
  107. >>> x = np.random.randn(1000)
  108. >>> detect = detect_non_silence(x)
  109. >>> assert x.shape == detect.shape
  110. >>> assert detect.dtype == np.bool
  111. """
  112. if x.shape[-1] < frame_length:
  113. return np.full(x.shape, fill_value=True, dtype=np.bool)
  114. if x.dtype.kind == "i":
  115. x = x.astype(np.float64)
  116. # framed_w: (C, T, F)
  117. framed_w = framing(
  118. x,
  119. frame_length=frame_length,
  120. frame_shift=frame_shift,
  121. centered=False,
  122. padded=True,
  123. )
  124. framed_w *= scipy.signal.get_window(window, frame_length).astype(framed_w.dtype)
  125. # power: (C, T)
  126. power = (framed_w ** 2).mean(axis=-1)
  127. # mean_power: (C, 1)
  128. mean_power = np.mean(power, axis=-1, keepdims=True)
  129. if np.all(mean_power == 0):
  130. return np.full(x.shape, fill_value=True, dtype=np.bool)
  131. # detect_frames: (C, T)
  132. detect_frames = power / mean_power > threshold
  133. # detects: (C, T, F)
  134. detects = np.broadcast_to(
  135. detect_frames[..., None], detect_frames.shape + (frame_shift,)
  136. )
  137. # detects: (C, TF)
  138. detects = detects.reshape(*detect_frames.shape[:-1], -1)
  139. # detects: (C, TF)
  140. return np.pad(
  141. detects,
  142. [(0, 0)] * (x.ndim - 1) + [(0, x.shape[-1] - detects.shape[-1])],
  143. mode="edge",
  144. )
  145. class CommonPreprocessor(AbsPreprocessor):
  146. def __init__(
  147. self,
  148. train: bool,
  149. token_type: str = None,
  150. token_list: Union[Path, str, Iterable[str]] = None,
  151. bpemodel: Union[Path, str, Iterable[str]] = None,
  152. text_cleaner: Collection[str] = None,
  153. g2p_type: str = None,
  154. unk_symbol: str = "<unk>",
  155. space_symbol: str = "<space>",
  156. non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
  157. delimiter: str = None,
  158. rir_scp: str = None,
  159. rir_apply_prob: float = 1.0,
  160. noise_scp: str = None,
  161. noise_apply_prob: float = 1.0,
  162. noise_db_range: str = "3_10",
  163. speech_volume_normalize: float = None,
  164. speech_name: str = "speech",
  165. text_name: str = "text",
  166. split_with_space: bool = False,
  167. seg_dict_file: str = None,
  168. ):
  169. super().__init__(train)
  170. self.train = train
  171. self.speech_name = speech_name
  172. self.text_name = text_name
  173. self.speech_volume_normalize = speech_volume_normalize
  174. self.rir_apply_prob = rir_apply_prob
  175. self.noise_apply_prob = noise_apply_prob
  176. self.split_with_space = split_with_space
  177. self.seg_dict = None
  178. if seg_dict_file is not None:
  179. self.seg_dict = {}
  180. with open(seg_dict_file) as f:
  181. lines = f.readlines()
  182. for line in lines:
  183. s = line.strip().split()
  184. key = s[0]
  185. value = s[1:]
  186. self.seg_dict[key] = " ".join(value)
  187. if token_type is not None:
  188. if token_list is None:
  189. raise ValueError("token_list is required if token_type is not None")
  190. self.text_cleaner = TextCleaner(text_cleaner)
  191. self.tokenizer = build_tokenizer(
  192. token_type=token_type,
  193. bpemodel=bpemodel,
  194. delimiter=delimiter,
  195. space_symbol=space_symbol,
  196. non_linguistic_symbols=non_linguistic_symbols,
  197. g2p_type=g2p_type,
  198. )
  199. self.token_id_converter = TokenIDConverter(
  200. token_list=token_list,
  201. unk_symbol=unk_symbol,
  202. )
  203. else:
  204. self.text_cleaner = None
  205. self.tokenizer = None
  206. self.token_id_converter = None
  207. if train and rir_scp is not None:
  208. self.rirs = []
  209. with open(rir_scp, "r", encoding="utf-8") as f:
  210. for line in f:
  211. sps = line.strip().split(None, 1)
  212. if len(sps) == 1:
  213. self.rirs.append(sps[0])
  214. else:
  215. self.rirs.append(sps[1])
  216. else:
  217. self.rirs = None
  218. if train and noise_scp is not None:
  219. self.noises = []
  220. with open(noise_scp, "r", encoding="utf-8") as f:
  221. for line in f:
  222. sps = line.strip().split(None, 1)
  223. if len(sps) == 1:
  224. self.noises.append(sps[0])
  225. else:
  226. self.noises.append(sps[1])
  227. sps = noise_db_range.split("_")
  228. if len(sps) == 1:
  229. self.noise_db_low, self.noise_db_high = float(sps[0])
  230. elif len(sps) == 2:
  231. self.noise_db_low, self.noise_db_high = float(sps[0]), float(sps[1])
  232. else:
  233. raise ValueError(
  234. "Format error: '{noise_db_range}' e.g. -3_4 -> [-3db,4db]"
  235. )
  236. else:
  237. self.noises = None
  238. def _speech_process(
  239. self, data: Dict[str, Union[str, np.ndarray]]
  240. ) -> Dict[str, Union[str, np.ndarray]]:
  241. assert check_argument_types()
  242. if self.speech_name in data:
  243. if self.train and (self.rirs is not None or self.noises is not None):
  244. speech = data[self.speech_name]
  245. nsamples = len(speech)
  246. # speech: (Nmic, Time)
  247. if speech.ndim == 1:
  248. speech = speech[None, :]
  249. else:
  250. speech = speech.T
  251. # Calc power on non shlence region
  252. power = (speech[detect_non_silence(speech)] ** 2).mean()
  253. # 1. Convolve RIR
  254. if self.rirs is not None and self.rir_apply_prob >= np.random.random():
  255. rir_path = np.random.choice(self.rirs)
  256. if rir_path is not None:
  257. rir, _ = soundfile.read(
  258. rir_path, dtype=np.float64, always_2d=True
  259. )
  260. # rir: (Nmic, Time)
  261. rir = rir.T
  262. # speech: (Nmic, Time)
  263. # Note that this operation doesn't change the signal length
  264. speech = scipy.signal.convolve(speech, rir, mode="full")[
  265. :, : speech.shape[1]
  266. ]
  267. # Reverse mean power to the original power
  268. power2 = (speech[detect_non_silence(speech)] ** 2).mean()
  269. speech = np.sqrt(power / max(power2, 1e-10)) * speech
  270. # 2. Add Noise
  271. if (
  272. self.noises is not None
  273. and self.noise_apply_prob >= np.random.random()
  274. ):
  275. noise_path = np.random.choice(self.noises)
  276. if noise_path is not None:
  277. noise_db = np.random.uniform(
  278. self.noise_db_low, self.noise_db_high
  279. )
  280. with soundfile.SoundFile(noise_path) as f:
  281. if f.frames == nsamples:
  282. noise = f.read(dtype=np.float64, always_2d=True)
  283. elif f.frames < nsamples:
  284. offset = np.random.randint(0, nsamples - f.frames)
  285. # noise: (Time, Nmic)
  286. noise = f.read(dtype=np.float64, always_2d=True)
  287. # Repeat noise
  288. noise = np.pad(
  289. noise,
  290. [(offset, nsamples - f.frames - offset), (0, 0)],
  291. mode="wrap",
  292. )
  293. else:
  294. offset = np.random.randint(0, f.frames - nsamples)
  295. f.seek(offset)
  296. # noise: (Time, Nmic)
  297. noise = f.read(
  298. nsamples, dtype=np.float64, always_2d=True
  299. )
  300. if len(noise) != nsamples:
  301. raise RuntimeError(f"Something wrong: {noise_path}")
  302. # noise: (Nmic, Time)
  303. noise = noise.T
  304. noise_power = (noise ** 2).mean()
  305. scale = (
  306. 10 ** (-noise_db / 20)
  307. * np.sqrt(power)
  308. / np.sqrt(max(noise_power, 1e-10))
  309. )
  310. speech = speech + scale * noise
  311. speech = speech.T
  312. ma = np.max(np.abs(speech))
  313. if ma > 1.0:
  314. speech /= ma
  315. data[self.speech_name] = speech
  316. if self.speech_volume_normalize is not None:
  317. speech = data[self.speech_name]
  318. ma = np.max(np.abs(speech))
  319. data[self.speech_name] = speech * self.speech_volume_normalize / ma
  320. assert check_return_type(data)
  321. return data
  322. def _text_process(
  323. self, data: Dict[str, Union[str, np.ndarray]]
  324. ) -> Dict[str, np.ndarray]:
  325. if self.text_name in data and self.tokenizer is not None:
  326. text = data[self.text_name]
  327. text = self.text_cleaner(text)
  328. if self.split_with_space:
  329. tokens = text.strip().split(" ")
  330. if self.seg_dict is not None:
  331. tokens = forward_segment("".join(tokens), self.seg_dict)
  332. tokens = seg_tokenize(tokens, self.seg_dict)
  333. else:
  334. tokens = self.tokenizer.text2tokens(text)
  335. text_ints = self.token_id_converter.tokens2ids(tokens)
  336. data[self.text_name] = np.array(text_ints, dtype=np.int64)
  337. assert check_return_type(data)
  338. return data
  339. def __call__(
  340. self, uid: str, data: Dict[str, Union[str, np.ndarray]]
  341. ) -> Dict[str, np.ndarray]:
  342. assert check_argument_types()
  343. data = self._speech_process(data)
  344. data = self._text_process(data)
  345. return data
  346. ## FIXME
  347. class LMPreprocessor(CommonPreprocessor):
  348. def __init__(
  349. self,
  350. train: bool,
  351. token_type: str = None,
  352. token_list: Union[Path, str, Iterable[str]] = None,
  353. bpemodel: Union[Path, str, Iterable[str]] = None,
  354. text_cleaner: Collection[str] = None,
  355. g2p_type: str = None,
  356. unk_symbol: str = "<unk>",
  357. space_symbol: str = "<space>",
  358. non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
  359. delimiter: str = None,
  360. rir_scp: str = None,
  361. rir_apply_prob: float = 1.0,
  362. noise_scp: str = None,
  363. noise_apply_prob: float = 1.0,
  364. noise_db_range: str = "3_10",
  365. speech_volume_normalize: float = None,
  366. speech_name: str = "speech",
  367. text_name: str = "text",
  368. split_with_space: bool = False,
  369. seg_dict_file: str = None,
  370. ):
  371. super().__init__(train,
  372. token_type,
  373. token_list,
  374. bpemodel,
  375. text_cleaner,
  376. g2p_type,
  377. unk_symbol,
  378. space_symbol,
  379. non_linguistic_symbols,
  380. delimiter,
  381. rir_scp,
  382. rir_apply_prob,
  383. noise_scp,
  384. noise_apply_prob,
  385. noise_db_range,
  386. speech_volume_normalize,
  387. speech_name,
  388. text_name,
  389. split_with_space,
  390. seg_dict_file,
  391. )
  392. def _text_process(
  393. self, data: Dict[str, Union[str, np.ndarray]]
  394. ) -> Dict[str, np.ndarray]:
  395. if self.text_name in data and self.tokenizer is not None:
  396. text = data[self.text_name]
  397. text = self.text_cleaner(text)
  398. if self.split_with_space:
  399. tokens = text.strip().split(" ")
  400. if self.seg_dict is not None:
  401. tokens = seg_tokenize_wo_pattern(tokens, self.seg_dict)
  402. else:
  403. tokens = self.tokenizer.text2tokens(text)
  404. text_ints = self.token_id_converter.tokens2ids(tokens)
  405. data[self.text_name] = np.array(text_ints, dtype=np.int64)
  406. assert check_return_type(data)
  407. return data
  408. class CommonPreprocessor_multi(AbsPreprocessor):
  409. def __init__(
  410. self,
  411. train: bool,
  412. token_type: str = None,
  413. token_list: Union[Path, str, Iterable[str]] = None,
  414. bpemodel: Union[Path, str, Iterable[str]] = None,
  415. text_cleaner: Collection[str] = None,
  416. g2p_type: str = None,
  417. unk_symbol: str = "<unk>",
  418. space_symbol: str = "<space>",
  419. non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
  420. delimiter: str = None,
  421. speech_name: str = "speech",
  422. text_name: List[str] = ["text"],
  423. ):
  424. super().__init__(train)
  425. self.train = train
  426. self.speech_name = speech_name
  427. self.text_name = text_name
  428. if token_type is not None:
  429. if token_list is None:
  430. raise ValueError("token_list is required if token_type is not None")
  431. self.text_cleaner = TextCleaner(text_cleaner)
  432. self.tokenizer = build_tokenizer(
  433. token_type=token_type,
  434. bpemodel=bpemodel,
  435. delimiter=delimiter,
  436. space_symbol=space_symbol,
  437. non_linguistic_symbols=non_linguistic_symbols,
  438. g2p_type=g2p_type,
  439. )
  440. self.token_id_converter = TokenIDConverter(
  441. token_list=token_list,
  442. unk_symbol=unk_symbol,
  443. )
  444. else:
  445. self.text_cleaner = None
  446. self.tokenizer = None
  447. self.token_id_converter = None
  448. def _text_process(
  449. self, data: Dict[str, Union[str, np.ndarray]]
  450. ) -> Dict[str, np.ndarray]:
  451. for text_n in self.text_name:
  452. if text_n in data and self.tokenizer is not None:
  453. text = data[text_n]
  454. text = self.text_cleaner(text)
  455. tokens = self.tokenizer.text2tokens(text)
  456. text_ints = self.token_id_converter.tokens2ids(tokens)
  457. data[text_n] = np.array(text_ints, dtype=np.int64)
  458. assert check_return_type(data)
  459. return data
  460. def __call__(
  461. self, uid: str, data: Dict[str, Union[str, np.ndarray]]
  462. ) -> Dict[str, np.ndarray]:
  463. assert check_argument_types()
  464. if self.speech_name in data:
  465. # Nothing now: candidates:
  466. # - STFT
  467. # - Fbank
  468. # - CMVN
  469. # - Data augmentation
  470. pass
  471. data = self._text_process(data)
  472. return data
  473. class MutliTokenizerCommonPreprocessor(CommonPreprocessor):
  474. def __init__(
  475. self,
  476. train: bool,
  477. token_type: List[str] = [None],
  478. token_list: List[Union[Path, str, Iterable[str]]] = [None],
  479. bpemodel: List[Union[Path, str, Iterable[str]]] = [None],
  480. text_cleaner: Collection[str] = None,
  481. g2p_type: str = None,
  482. unk_symbol: str = "<unk>",
  483. space_symbol: str = "<space>",
  484. non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
  485. delimiter: str = None,
  486. rir_scp: str = None,
  487. rir_apply_prob: float = 1.0,
  488. noise_scp: str = None,
  489. noise_apply_prob: float = 1.0,
  490. noise_db_range: str = "3_10",
  491. speech_volume_normalize: float = None,
  492. speech_name: str = "speech",
  493. text_name: List[str] = ["text"],
  494. ):
  495. # TODO(jiatong): sync with Kamo and Jing on interface for preprocessor
  496. super().__init__(
  497. train=train,
  498. token_type=token_type[0],
  499. token_list=token_list[0],
  500. bpemodel=bpemodel[0],
  501. text_cleaner=text_cleaner,
  502. g2p_type=g2p_type,
  503. unk_symbol=unk_symbol,
  504. space_symbol=space_symbol,
  505. non_linguistic_symbols=non_linguistic_symbols,
  506. delimiter=delimiter,
  507. speech_name=speech_name,
  508. text_name=text_name[0],
  509. rir_scp=rir_scp,
  510. rir_apply_prob=rir_apply_prob,
  511. noise_scp=noise_scp,
  512. noise_apply_prob=noise_apply_prob,
  513. noise_db_range=noise_db_range,
  514. speech_volume_normalize=speech_volume_normalize,
  515. )
  516. assert (
  517. len(token_type) == len(token_list) == len(bpemodel) == len(text_name)
  518. ), "token_type, token_list, bpemodel, or processing text_name mismatched"
  519. self.num_tokenizer = len(token_type)
  520. self.tokenizer = []
  521. self.token_id_converter = []
  522. for i in range(self.num_tokenizer):
  523. if token_type[i] is not None:
  524. if token_list[i] is None:
  525. raise ValueError("token_list is required if token_type is not None")
  526. self.tokenizer.append(
  527. build_tokenizer(
  528. token_type=token_type[i],
  529. bpemodel=bpemodel[i],
  530. delimiter=delimiter,
  531. space_symbol=space_symbol,
  532. non_linguistic_symbols=non_linguistic_symbols,
  533. g2p_type=g2p_type,
  534. )
  535. )
  536. self.token_id_converter.append(
  537. TokenIDConverter(
  538. token_list=token_list[i],
  539. unk_symbol=unk_symbol,
  540. )
  541. )
  542. else:
  543. self.tokenizer.append(None)
  544. self.token_id_converter.append(None)
  545. self.text_cleaner = TextCleaner(text_cleaner)
  546. self.text_name = text_name # override the text_name from CommonPreprocessor
  547. def _text_process(
  548. self, data: Dict[str, Union[str, np.ndarray]]
  549. ) -> Dict[str, np.ndarray]:
  550. for i in range(self.num_tokenizer):
  551. text_name = self.text_name[i]
  552. if text_name in data and self.tokenizer[i] is not None:
  553. text = data[text_name]
  554. text = self.text_cleaner(text)
  555. tokens = self.tokenizer[i].text2tokens(text)
  556. text_ints = self.token_id_converter[i].tokens2ids(tokens)
  557. data[text_name] = np.array(text_ints, dtype=np.int64)
  558. assert check_return_type(data)
  559. return data
  560. class CodeMixTokenizerCommonPreprocessor(CommonPreprocessor):
  561. def __init__(
  562. self,
  563. train: bool,
  564. token_type: str = None,
  565. token_list: Union[Path, str, Iterable[str]] = None,
  566. bpemodel: Union[Path, str, Iterable[str]] = None,
  567. text_cleaner: Collection[str] = None,
  568. g2p_type: str = None,
  569. unk_symbol: str = "<unk>",
  570. space_symbol: str = "<space>",
  571. non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
  572. delimiter: str = None,
  573. rir_scp: str = None,
  574. rir_apply_prob: float = 1.0,
  575. noise_scp: str = None,
  576. noise_apply_prob: float = 1.0,
  577. noise_db_range: str = "3_10",
  578. speech_volume_normalize: float = None,
  579. speech_name: str = "speech",
  580. text_name: str = "text",
  581. split_text_name: str = "split_text",
  582. split_with_space: bool = False,
  583. seg_dict_file: str = None,
  584. ):
  585. super().__init__(
  586. train=train,
  587. # Force to use word.
  588. token_type="word",
  589. token_list=token_list,
  590. bpemodel=bpemodel,
  591. text_cleaner=text_cleaner,
  592. g2p_type=g2p_type,
  593. unk_symbol=unk_symbol,
  594. space_symbol=space_symbol,
  595. non_linguistic_symbols=non_linguistic_symbols,
  596. delimiter=delimiter,
  597. speech_name=speech_name,
  598. text_name=text_name,
  599. rir_scp=rir_scp,
  600. rir_apply_prob=rir_apply_prob,
  601. noise_scp=noise_scp,
  602. noise_apply_prob=noise_apply_prob,
  603. noise_db_range=noise_db_range,
  604. speech_volume_normalize=speech_volume_normalize,
  605. split_with_space=split_with_space,
  606. seg_dict_file=seg_dict_file,
  607. )
  608. # The data field name for split text.
  609. self.split_text_name = split_text_name
  610. @classmethod
  611. def split_words(cls, text: str):
  612. words = []
  613. segs = text.split()
  614. for seg in segs:
  615. # There is no space in seg.
  616. current_word = ""
  617. for c in seg:
  618. if len(c.encode()) == 1:
  619. # This is an ASCII char.
  620. current_word += c
  621. else:
  622. # This is a Chinese char.
  623. if len(current_word) > 0:
  624. words.append(current_word)
  625. current_word = ""
  626. words.append(c)
  627. if len(current_word) > 0:
  628. words.append(current_word)
  629. return words
  630. def __call__(
  631. self, uid: str, data: Dict[str, Union[list, str, np.ndarray]]
  632. ) -> Dict[str, Union[list, np.ndarray]]:
  633. assert check_argument_types()
  634. # Split words.
  635. if isinstance(data[self.text_name], str):
  636. split_text = self.split_words(data[self.text_name])
  637. else:
  638. split_text = data[self.text_name]
  639. data[self.text_name] = " ".join(split_text)
  640. data = self._speech_process(data)
  641. data = self._text_process(data)
  642. data[self.split_text_name] = split_text
  643. return data
  644. def pop_split_text_data(self, data: Dict[str, Union[str, np.ndarray]]):
  645. result = data[self.split_text_name]
  646. del data[self.split_text_name]
  647. return result
  648. class PuncTrainTokenizerCommonPreprocessor(CommonPreprocessor):
  649. def __init__(
  650. self,
  651. train: bool,
  652. token_type: List[str] = [None],
  653. token_list: List[Union[Path, str, Iterable[str]]] = [None],
  654. bpemodel: List[Union[Path, str, Iterable[str]]] = [None],
  655. text_cleaner: Collection[str] = None,
  656. g2p_type: str = None,
  657. unk_symbol: str = "<unk>",
  658. space_symbol: str = "<space>",
  659. non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
  660. delimiter: str = None,
  661. rir_scp: str = None,
  662. rir_apply_prob: float = 1.0,
  663. noise_scp: str = None,
  664. noise_apply_prob: float = 1.0,
  665. noise_db_range: str = "3_10",
  666. speech_volume_normalize: float = None,
  667. speech_name: str = "speech",
  668. text_name: List[str] = ["text"],
  669. vad_name: str = "vad_indexes",
  670. ):
  671. # TODO(jiatong): sync with Kamo and Jing on interface for preprocessor
  672. super().__init__(
  673. train=train,
  674. token_type=token_type[0],
  675. token_list=token_list[0],
  676. bpemodel=bpemodel[0],
  677. text_cleaner=text_cleaner,
  678. g2p_type=g2p_type,
  679. unk_symbol=unk_symbol,
  680. space_symbol=space_symbol,
  681. non_linguistic_symbols=non_linguistic_symbols,
  682. delimiter=delimiter,
  683. speech_name=speech_name,
  684. text_name=text_name[0],
  685. rir_scp=rir_scp,
  686. rir_apply_prob=rir_apply_prob,
  687. noise_scp=noise_scp,
  688. noise_apply_prob=noise_apply_prob,
  689. noise_db_range=noise_db_range,
  690. speech_volume_normalize=speech_volume_normalize,
  691. )
  692. assert (
  693. len(token_type) == len(token_list) == len(bpemodel) == len(text_name)
  694. ), "token_type, token_list, bpemodel, or processing text_name mismatched"
  695. self.num_tokenizer = len(token_type)
  696. self.tokenizer = []
  697. self.token_id_converter = []
  698. for i in range(self.num_tokenizer):
  699. if token_type[i] is not None:
  700. if token_list[i] is None:
  701. raise ValueError("token_list is required if token_type is not None")
  702. self.tokenizer.append(
  703. build_tokenizer(
  704. token_type=token_type[i],
  705. bpemodel=bpemodel[i],
  706. delimiter=delimiter,
  707. space_symbol=space_symbol,
  708. non_linguistic_symbols=non_linguistic_symbols,
  709. g2p_type=g2p_type,
  710. )
  711. )
  712. self.token_id_converter.append(
  713. TokenIDConverter(
  714. token_list=token_list[i],
  715. unk_symbol=unk_symbol,
  716. )
  717. )
  718. else:
  719. self.tokenizer.append(None)
  720. self.token_id_converter.append(None)
  721. self.text_cleaner = TextCleaner(text_cleaner)
  722. self.text_name = text_name # override the text_name from CommonPreprocessor
  723. self.vad_name = vad_name
  724. def _text_process(
  725. self, data: Dict[str, Union[str, np.ndarray]]
  726. ) -> Dict[str, np.ndarray]:
  727. for i in range(self.num_tokenizer):
  728. text_name = self.text_name[i]
  729. if text_name in data and self.tokenizer[i] is not None:
  730. text = data[text_name]
  731. text = self.text_cleaner(text)
  732. tokens = self.tokenizer[i].text2tokens(text)
  733. if "vad:" in tokens[-1]:
  734. vad = tokens[-1][4:]
  735. tokens = tokens[:-1]
  736. if len(vad) == 0:
  737. vad = -1
  738. else:
  739. vad = int(vad)
  740. data[self.vad_name] = np.array([vad], dtype=np.int64)
  741. text_ints = self.token_id_converter[i].tokens2ids(tokens)
  742. data[text_name] = np.array(text_ints, dtype=np.int64)