preprocessor.py 32 KB

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