wehub-resource-sync
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import logging
|
|
import concurrent.futures
|
|
import librosa
|
|
import torch.distributed as dist
|
|
from typing import Collection
|
|
import torch
|
|
import torchaudio
|
|
from torch import nn
|
|
import random
|
|
import re
|
|
from funasr.tokenizer.cleaner import TextCleaner
|
|
from funasr.register import tables
|
|
|
|
|
|
@tables.register("preprocessor_classes", "SpeechPreprocessSpeedPerturb")
|
|
class SpeechPreprocessSpeedPerturb(nn.Module):
|
|
def __init__(self, speed_perturb: list = None, **kwargs):
|
|
"""Initialize SpeechPreprocessSpeedPerturb.
|
|
|
|
Args:
|
|
speed_perturb: TODO.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
super().__init__()
|
|
self.speed_perturb = speed_perturb
|
|
|
|
def forward(self, waveform, fs, **kwargs):
|
|
"""Forward pass for training.
|
|
|
|
Args:
|
|
waveform: TODO.
|
|
fs: TODO.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
if self.speed_perturb is None:
|
|
return waveform
|
|
speed = random.choice(self.speed_perturb)
|
|
if speed != 1.0:
|
|
if not isinstance(waveform, torch.Tensor):
|
|
waveform = torch.tensor(waveform)
|
|
waveform, _ = torchaudio.sox_effects.apply_effects_tensor(
|
|
waveform.view(1, -1), fs, [["speed", str(speed)], ["rate", str(fs)]]
|
|
)
|
|
waveform = waveform.view(-1)
|
|
|
|
return waveform
|
|
|
|
|
|
@tables.register("preprocessor_classes", "TextPreprocessSegDict")
|
|
class TextPreprocessSegDict(nn.Module):
|
|
def __init__(
|
|
self,
|
|
seg_dict: str = None,
|
|
text_cleaner: Collection[str] = None,
|
|
split_with_space: bool = False,
|
|
**kwargs
|
|
):
|
|
"""Initialize TextPreprocessSegDict.
|
|
|
|
Args:
|
|
seg_dict: TODO.
|
|
text_cleaner: TODO.
|
|
split_with_space: TODO.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
super().__init__()
|
|
|
|
self.text_cleaner = TextCleaner(text_cleaner)
|
|
|
|
def forward(self, text, **kwargs):
|
|
"""Forward pass for training.
|
|
|
|
Args:
|
|
text: Text tensor or string input.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
text = self.text_cleaner(text)
|
|
|
|
return text
|