python/FunASR-XL.git

			@@ -38,6 +38,7 @@
			from funasr.models.decoder.transformer_decoder import ParaformerDecoderSAN
			from funasr.models.decoder.transformer_decoder import TransformerDecoder
			from funasr.models.decoder.contextual_decoder import ContextualParaformerDecoder
			from funasr.models.decoder.transformer_decoder import SAAsrTransformerDecoder
			from funasr.models.e2e_asr import ASRModel
			from funasr.models.decoder.rnnt_decoder import RNNTDecoder
			from funasr.models.joint_net.joint_network import JointNetwork
			@@ -45,6 +46,7 @@
			from funasr.models.e2e_asr_contextual_paraformer import NeatContextualParaformer
			from funasr.models.e2e_tp import TimestampPredictor
			from funasr.models.e2e_asr_mfcca import MFCCA
			from funasr.models.e2e_sa_asr import SAASRModel
			from funasr.models.e2e_uni_asr import UniASR
			from funasr.models.e2e_asr_transducer import TransducerModel, UnifiedTransducerModel
			from funasr.models.encoder.abs_encoder import AbsEncoder
			@@ -54,6 +56,7 @@
			from funasr.models.encoder.sanm_encoder import SANMEncoder, SANMEncoderChunkOpt
			from funasr.models.encoder.transformer_encoder import TransformerEncoder
			from funasr.models.encoder.mfcca_encoder import MFCCAEncoder
			from funasr.models.encoder.resnet34_encoder import ResNet34Diar
			from funasr.models.frontend.abs_frontend import AbsFrontend
			from funasr.models.frontend.default import DefaultFrontend
			from funasr.models.frontend.default import MultiChannelFrontend
			@@ -132,6 +135,9 @@
			neatcontextual_paraformer=NeatContextualParaformer,
			mfcca=MFCCA,
			timestamp_prediction=TimestampPredictor,
			rnnt=TransducerModel,
			rnnt_unified=UnifiedTransducerModel,
			sa_asr=SAASRModel,
			),
			type_check=FunASRModel,
			default="asr",
			@@ -173,6 +179,27 @@
			type_check=AbsEncoder,
			default="rnn",
			)
			asr_encoder_choices = ClassChoices(
			"asr_encoder",
			classes=dict(
			conformer=ConformerEncoder,
			transformer=TransformerEncoder,
			rnn=RNNEncoder,
			sanm=SANMEncoder,
			sanm_chunk_opt=SANMEncoderChunkOpt,
			data2vec_encoder=Data2VecEncoder,
			mfcca_enc=MFCCAEncoder,
			),
			type_check=AbsEncoder,
			default="rnn",
			)
			spk_encoder_choices = ClassChoices(
			"spk_encoder",
			classes=dict(
			resnet34_diar=ResNet34Diar,
			),
			default="resnet34_diar",
			)
			postencoder_choices = ClassChoices(
			name="postencoder",
			classes=dict(
			@@ -195,6 +222,7 @@
			paraformer_decoder_sanm=ParaformerSANMDecoder,
			paraformer_decoder_san=ParaformerDecoderSAN,
			contextual_paraformer_decoder=ContextualParaformerDecoder,
			sa_decoder=SAAsrTransformerDecoder,
			),
			type_check=AbsDecoder,
			default="rnn",
			@@ -222,6 +250,15 @@
			),
			type_check=RNNTDecoder,
			default="rnnt",
			)

			joint_network_choices = ClassChoices(
			name="joint_network",
			classes=dict(
			joint_network=JointNetwork,
			),
			default="joint_network",
			optional=True,
			)

			predictor_choices = ClassChoices(
			@@ -290,6 +327,8 @@
			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
			@@ -315,6 +354,12 @@
			type=str2bool,
			default=True,
			help="whether to split text using <space>",
			)
			group.add_argument(
			"--max_spk_num",
			type=int_or_none,
			default=None,
			help="A text mapping int-id to token",
			)
			group.add_argument(
			"--seg_dict_file",
			@@ -349,12 +394,6 @@
			action=NestedDictAction,
			default=get_default_kwargs(CTC),
			help="The keyword arguments for CTC class.",
			)
			group.add_argument(
			"--joint_net_conf",
			action=NestedDictAction,
			default=None,
			help="The keyword arguments for joint network class.",
			)

			group = parser.add_argument_group(description="Preprocess related")
			@@ -1360,254 +1399,22 @@
			return retval


			class ASRTransducerTask(AbsTask):
			class ASRTransducerTask(ASRTask):
			"""ASR Transducer Task definition."""

			num_optimizers: int = 1

			class_choices_list = [
			model_choices,
			frontend_choices,
			specaug_choices,
			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 <space>",
			)
			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 --<name> and --<name>_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:
			@@ -1675,7 +1482,7 @@
			decoder_output_size = decoder.output_size

			if getattr(args, "decoder", None) is not None:
			att_decoder_class = decoder_choices.get_class(args.att_decoder)
			att_decoder_class = decoder_choices.get_class(args.decoder)

			att_decoder = att_decoder_class(
			vocab_size=vocab_size,
			@@ -1693,35 +1500,23 @@
			)

			# 7. Build model
			try:
			model_class = model_choices.get_class(args.model)
			except AttributeError:
			model_class = model_choices.get_class("rnnt_unified")

			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,
			)

			model = model_class(
			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,
			)
			# 8. Initialize model
			if args.init is not None:
			raise NotImplementedError(
			@@ -1732,3 +1527,123 @@
			#assert check_return_type(model)

			return model


			class ASRTaskSAASR(ASRTask):
			# If you need more than one optimizers, change this value
			num_optimizers: int = 1

			# Add variable objects configurations
			class_choices_list = [
			# --frontend and --frontend_conf
			frontend_choices,
			# --specaug and --specaug_conf
			specaug_choices,
			# --normalize and --normalize_conf
			normalize_choices,
			# --model and --model_conf
			model_choices,
			# --preencoder and --preencoder_conf
			preencoder_choices,
			# --encoder and --encoder_conf
			# --asr_encoder and --asr_encoder_conf
			asr_encoder_choices,
			# --spk_encoder and --spk_encoder_conf
			spk_encoder_choices,
			# --decoder and --decoder_conf
			decoder_choices,
			]

			# If you need to modify train() or eval() procedures, change Trainer class here
			trainer = Trainer

			@classmethod
			def build_model(cls, args: argparse.Namespace):
			assert check_argument_types()
			if isinstance(args.token_list, str):
			with open(args.token_list, encoding="utf-8") as f:
			token_list = [line.rstrip() for line in f]

			# Overwriting token_list to keep it as "portable".
			args.token_list = list(token_list)
			elif isinstance(args.token_list, (tuple, list)):
			token_list = list(args.token_list)
			else:
			raise RuntimeError("token_list must be str or list")
			vocab_size = len(token_list)
			logging.info(f"Vocabulary size: {vocab_size}")

			# 1. frontend
			if args.input_size is None:
			# Extract features in the model
			frontend_class = frontend_choices.get_class(args.frontend)
			if args.frontend == 'wav_frontend' or args.frontend == "multichannelfrontend":
			frontend = frontend_class(cmvn_file=args.cmvn_file, **args.frontend_conf)
			else:
			frontend = frontend_class(**args.frontend_conf)
			input_size = frontend.output_size()
			else:
			# Give features from data-loader
			args.frontend = None
			args.frontend_conf = {}
			frontend = None
			input_size = args.input_size

			# 2. Data augmentation for spectrogram
			if args.specaug is not None:
			specaug_class = specaug_choices.get_class(args.specaug)
			specaug = specaug_class(**args.specaug_conf)
			else:
			specaug = None

			# 3. Normalization layer
			if args.normalize is not None:
			normalize_class = normalize_choices.get_class(args.normalize)
			normalize = normalize_class(**args.normalize_conf)
			else:
			normalize = None

			# 5. Encoder
			asr_encoder_class = asr_encoder_choices.get_class(args.asr_encoder)
			asr_encoder = asr_encoder_class(input_size=input_size, **args.asr_encoder_conf)
			spk_encoder_class = spk_encoder_choices.get_class(args.spk_encoder)
			spk_encoder = spk_encoder_class(input_size=input_size, **args.spk_encoder_conf)

			# 7. Decoder
			decoder_class = decoder_choices.get_class(args.decoder)
			decoder = decoder_class(
			vocab_size=vocab_size,
			encoder_output_size=asr_encoder.output_size(),
			**args.decoder_conf,
			)

			# 8. CTC
			ctc = CTC(
			odim=vocab_size, encoder_output_size=asr_encoder.output_size(), **args.ctc_conf
			)

			# import ipdb;ipdb.set_trace()
			# 9. Build model
			try:
			model_class = model_choices.get_class(args.model)
			except AttributeError:
			model_class = model_choices.get_class("asr")
			model = model_class(
			vocab_size=vocab_size,
			frontend=frontend,
			specaug=specaug,
			normalize=normalize,
			asr_encoder=asr_encoder,
			spk_encoder=spk_encoder,
			decoder=decoder,
			ctc=ctc,
			token_list=token_list,
			**args.model_conf,
			)

			# 10. Initialize
			if args.init is not None:
			initialize(model, args.init)

			assert check_return_type(model)
			return model