From 9d01231fa69672fc7c9b4bf81ef466bb0189788c Mon Sep 17 00:00:00 2001 From: aky15 Date: Wed, 17 May 2023 17:34:21 +0800 Subject: [PATCH] =?UTF-8?q?rnnt=E7=BB=A7=E6=89=BFASRTask?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- funasr/bin/asr_train.py | 2 + funasr/build_utils/build_asr_model.py | 86 +++++++- funasr/models/encoder/conformer_encoder.py | 2 +- funasr/tasks/asr.py | 239 +-------------------- 4 files changed, 91 insertions(+), 238 deletions(-) diff --git a/funasr/bin/asr_train.py b/funasr/bin/asr_train.py index 38a42b3e3..fd973a41b 100755 --- a/funasr/bin/asr_train.py +++ b/funasr/bin/asr_train.py @@ -36,6 +36,8 @@ def main(args=None, cmd=None): from funasr.tasks.asr import ASRTaskParaformer as ASRTask if args.mode == "uniasr": from funasr.tasks.asr import ASRTaskUniASR as ASRTask + if args.mode == "rnnt": + from funasr.tasks.asr import ASRTransducerTask as ASRTask ASRTask.main(args=args, cmd=cmd) diff --git a/funasr/build_utils/build_asr_model.py b/funasr/build_utils/build_asr_model.py index d8cbba572..718736b9c 100644 --- a/funasr/build_utils/build_asr_model.py +++ b/funasr/build_utils/build_asr_model.py @@ -19,12 +19,15 @@ from funasr.models.decoder.transformer_decoder import ( ) from funasr.models.decoder.transformer_decoder import ParaformerDecoderSAN from funasr.models.decoder.transformer_decoder import TransformerDecoder +from funasr.models.decoder.rnnt_decoder import RNNTDecoder +from funasr.models.joint_net.joint_network import JointNetwork from funasr.models.e2e_asr import ASRModel from funasr.models.e2e_asr_mfcca import MFCCA from funasr.models.e2e_asr_paraformer import Paraformer, ParaformerBert, BiCifParaformer, ContextualParaformer from funasr.models.e2e_tp import TimestampPredictor from funasr.models.e2e_uni_asr import UniASR -from funasr.models.encoder.conformer_encoder import ConformerEncoder +from funasr.models.e2e_asr_transducer import TransducerModel, UnifiedTransducerModel +from funasr.models.encoder.conformer_encoder import ConformerEncoder, ConformerChunkEncoder from funasr.models.encoder.data2vec_encoder import Data2VecEncoder from funasr.models.encoder.mfcca_encoder import MFCCAEncoder from funasr.models.encoder.rnn_encoder import RNNEncoder @@ -97,6 +100,7 @@ encoder_choices = ClassChoices( sanm_chunk_opt=SANMEncoderChunkOpt, data2vec_encoder=Data2VecEncoder, mfcca_enc=MFCCAEncoder, + chunk_conformer=ConformerChunkEncoder, ), default="rnn", ) @@ -171,6 +175,23 @@ stride_conv_choices = ClassChoices( default="stride_conv1d", optional=True, ) +rnnt_decoder_choices = ClassChoices( + name="rnnt_decoder", + classes=dict( + rnnt=RNNTDecoder, + ), + default="rnnt", + optional=True, +) +joint_network_choices = ClassChoices( + name="joint_network", + classes=dict( + joint_network=JointNetwork, + ), + default="joint_network", + optional=True, +) + class_choices_list = [ # --frontend and --frontend_conf frontend_choices, @@ -194,6 +215,10 @@ class_choices_list = [ predictor_choices2, # --stride_conv and --stride_conv_conf stride_conv_choices, + # --rnnt_decoder and --rnnt_decoder_conf + rnnt_decoder_choices, + # --joint_network and --joint_network_conf + joint_network_choices, ] @@ -342,6 +367,63 @@ def build_asr_model(args): token_list=token_list, **args.model_conf, ) + elif args.model == "rnnt": + # 5. Decoder + encoder_output_size = encoder.output_size() + + rnnt_decoder_class = rnnt_decoder_choices.get_class(args.rnnt_decoder) + decoder = rnnt_decoder_class( + vocab_size, + **args.rnnt_decoder_conf, + ) + decoder_output_size = decoder.output_size + + if getattr(args, "decoder", None) is not None: + att_decoder_class = decoder_choices.get_class(args.decoder) + + att_decoder = att_decoder_class( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + **args.decoder_conf, + ) + else: + att_decoder = None + # 6. Joint Network + joint_network = JointNetwork( + vocab_size, + encoder_output_size, + decoder_output_size, + **args.joint_network_conf, + ) + + # 7. Build model + if hasattr(encoder, 'unified_model_training') and encoder.unified_model_training: + model = UnifiedTransducerModel( + vocab_size=vocab_size, + token_list=token_list, + frontend=frontend, + specaug=specaug, + normalize=normalize, + encoder=encoder, + decoder=decoder, + att_decoder=att_decoder, + joint_network=joint_network, + **args.model_conf, + ) + + else: + model = TransducerModel( + vocab_size=vocab_size, + token_list=token_list, + frontend=frontend, + specaug=specaug, + normalize=normalize, + encoder=encoder, + decoder=decoder, + att_decoder=att_decoder, + joint_network=joint_network, + **args.model_conf, + ) else: raise NotImplementedError("Not supported model: {}".format(args.model)) @@ -349,4 +431,4 @@ def build_asr_model(args): if args.init is not None: initialize(model, args.init) - return model \ No newline at end of file + return model diff --git a/funasr/models/encoder/conformer_encoder.py b/funasr/models/encoder/conformer_encoder.py index aa3b67ecc..5f20deec4 100644 --- a/funasr/models/encoder/conformer_encoder.py +++ b/funasr/models/encoder/conformer_encoder.py @@ -1078,7 +1078,7 @@ class ConformerChunkEncoder(AbsEncoder): limit_size, ) - mask = make_source_mask(x_len) + mask = make_source_mask(x_len).to(x.device) if self.unified_model_training: chunk_size = self.default_chunk_size + torch.randint(-self.jitter_range, self.jitter_range+1, (1,)).item() diff --git a/funasr/tasks/asr.py b/funasr/tasks/asr.py index d218902a8..0bb056365 100644 --- a/funasr/tasks/asr.py +++ b/funasr/tasks/asr.py @@ -290,6 +290,8 @@ class ASRTask(AbsTask): predictor_choices2, # --stride_conv and --stride_conv_conf stride_conv_choices, + # --rnnt_decoder and --rnnt_decoder_conf + rnnt_decoder_choices, ] # If you need to modify train() or eval() procedures, change Trainer class here @@ -1360,7 +1362,7 @@ class ASRTaskAligner(ASRTaskParaformer): return retval -class ASRTransducerTask(AbsTask): +class ASRTransducerTask(ASRTask): """ASR Transducer Task definition.""" num_optimizers: int = 1 @@ -1371,244 +1373,11 @@ class ASRTransducerTask(AbsTask): normalize_choices, encoder_choices, rnnt_decoder_choices, + joint_network_choices, ] trainer = Trainer - @classmethod - def add_task_arguments(cls, parser: argparse.ArgumentParser): - """Add Transducer task arguments. - Args: - cls: ASRTransducerTask object. - parser: Transducer arguments parser. - """ - group = parser.add_argument_group(description="Task related.") - - # required = parser.get_default("required") - # required += ["token_list"] - - group.add_argument( - "--token_list", - type=str_or_none, - default=None, - help="Integer-string mapper for tokens.", - ) - group.add_argument( - "--split_with_space", - type=str2bool, - default=True, - help="whether to split text using ", - ) - group.add_argument( - "--input_size", - type=int_or_none, - default=None, - help="The number of dimensions for input features.", - ) - group.add_argument( - "--init", - type=str_or_none, - default=None, - help="Type of model initialization to use.", - ) - group.add_argument( - "--model_conf", - action=NestedDictAction, - default=get_default_kwargs(TransducerModel), - help="The keyword arguments for the model class.", - ) - # group.add_argument( - # "--encoder_conf", - # action=NestedDictAction, - # default={}, - # help="The keyword arguments for the encoder class.", - # ) - group.add_argument( - "--joint_network_conf", - action=NestedDictAction, - default={}, - help="The keyword arguments for the joint network class.", - ) - group = parser.add_argument_group(description="Preprocess related.") - group.add_argument( - "--use_preprocessor", - type=str2bool, - default=True, - help="Whether to apply preprocessing to input data.", - ) - group.add_argument( - "--token_type", - type=str, - default="bpe", - choices=["bpe", "char", "word", "phn"], - help="The type of tokens to use during tokenization.", - ) - group.add_argument( - "--bpemodel", - type=str_or_none, - default=None, - help="The path of the sentencepiece model.", - ) - parser.add_argument( - "--non_linguistic_symbols", - type=str_or_none, - help="The 'non_linguistic_symbols' file path.", - ) - parser.add_argument( - "--cleaner", - type=str_or_none, - choices=[None, "tacotron", "jaconv", "vietnamese"], - default=None, - help="Text cleaner to use.", - ) - parser.add_argument( - "--g2p", - type=str_or_none, - choices=g2p_choices, - default=None, - help="g2p method to use if --token_type=phn.", - ) - parser.add_argument( - "--speech_volume_normalize", - type=float_or_none, - default=None, - help="Normalization value for maximum amplitude scaling.", - ) - parser.add_argument( - "--rir_scp", - type=str_or_none, - default=None, - help="The RIR SCP file path.", - ) - parser.add_argument( - "--rir_apply_prob", - type=float, - default=1.0, - help="The probability of the applied RIR convolution.", - ) - parser.add_argument( - "--noise_scp", - type=str_or_none, - default=None, - help="The path of noise SCP file.", - ) - parser.add_argument( - "--noise_apply_prob", - type=float, - default=1.0, - help="The probability of the applied noise addition.", - ) - parser.add_argument( - "--noise_db_range", - type=str, - default="13_15", - help="The range of the noise decibel level.", - ) - for class_choices in cls.class_choices_list: - # Append -- and --_conf. - # e.g. --decoder and --decoder_conf - class_choices.add_arguments(group) - - @classmethod - def build_collate_fn( - cls, args: argparse.Namespace, train: bool - ) -> Callable[ - [Collection[Tuple[str, Dict[str, np.ndarray]]]], - Tuple[List[str], Dict[str, torch.Tensor]], - ]: - """Build collate function. - Args: - cls: ASRTransducerTask object. - args: Task arguments. - train: Training mode. - Return: - : Callable collate function. - """ - assert check_argument_types() - - return CommonCollateFn(float_pad_value=0.0, int_pad_value=-1) - - @classmethod - def build_preprocess_fn( - cls, args: argparse.Namespace, train: bool - ) -> Optional[Callable[[str, Dict[str, np.array]], Dict[str, np.ndarray]]]: - """Build pre-processing function. - Args: - cls: ASRTransducerTask object. - args: Task arguments. - train: Training mode. - Return: - : Callable pre-processing function. - """ - assert check_argument_types() - - if args.use_preprocessor: - retval = CommonPreprocessor( - train=train, - token_type=args.token_type, - token_list=args.token_list, - bpemodel=args.bpemodel, - non_linguistic_symbols=args.non_linguistic_symbols, - text_cleaner=args.cleaner, - g2p_type=args.g2p, - split_with_space=args.split_with_space if hasattr(args, "split_with_space") else False, - rir_scp=args.rir_scp if hasattr(args, "rir_scp") else None, - rir_apply_prob=args.rir_apply_prob - if hasattr(args, "rir_apply_prob") - else 1.0, - noise_scp=args.noise_scp if hasattr(args, "noise_scp") else None, - noise_apply_prob=args.noise_apply_prob - if hasattr(args, "noise_apply_prob") - else 1.0, - noise_db_range=args.noise_db_range - if hasattr(args, "noise_db_range") - else "13_15", - speech_volume_normalize=args.speech_volume_normalize - if hasattr(args, "rir_scp") - else None, - ) - else: - retval = None - - assert check_return_type(retval) - return retval - - @classmethod - def required_data_names( - cls, train: bool = True, inference: bool = False - ) -> Tuple[str, ...]: - """Required data depending on task mode. - Args: - cls: ASRTransducerTask object. - train: Training mode. - inference: Inference mode. - Return: - retval: Required task data. - """ - if not inference: - retval = ("speech", "text") - else: - retval = ("speech",) - - return retval - - @classmethod - def optional_data_names( - cls, train: bool = True, inference: bool = False - ) -> Tuple[str, ...]: - """Optional data depending on task mode. - Args: - cls: ASRTransducerTask object. - train: Training mode. - inference: Inference mode. - Return: - retval: Optional task data. - """ - retval = () - assert check_return_type(retval) - - return retval - @classmethod def build_model(cls, args: argparse.Namespace) -> TransducerModel: """Required data depending on task mode.