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(
			@@ -319,6 +356,12 @@
			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",
			type=str,
			default=None,
			@@ -351,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")
			@@ -1368,6 +1405,7 @@
			num_optimizers: int = 1

			class_choices_list = [
			model_choices,
			frontend_choices,
			specaug_choices,
			normalize_choices,
			@@ -1444,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,
			@@ -1462,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(
			@@ -1501,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