python/FunASR-XL.git

			@@ -6,7 +6,6 @@
			from funasr.models.decoder.abs_decoder import AbsDecoder
			from funasr.models.decoder.contextual_decoder import ContextualParaformerDecoder
			from funasr.models.decoder.rnn_decoder import RNNDecoder
			from funasr.models.decoder.rnnt_decoder import RNNTDecoder
			from funasr.models.decoder.sanm_decoder import ParaformerSANMDecoder, FsmnDecoderSCAMAOpt
			from funasr.models.decoder.transformer_decoder import (
			DynamicConvolution2DTransformerDecoder, # noqa: H301
			@@ -20,19 +19,30 @@
			)
			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.decoder.transformer_decoder import SAAsrTransformerDecoder
			from funasr.models.e2e_asr import ASRModel
			from funasr.models.e2e_asr_contextual_paraformer import NeatContextualParaformer
			from funasr.models.e2e_asr_mfcca import MFCCA
			from funasr.models.e2e_asr_paraformer import Paraformer, ParaformerOnline, ParaformerBert, BiCifParaformer, \
			ContextualParaformer

			from funasr.models.e2e_asr_transducer import TransducerModel, UnifiedTransducerModel
			from funasr.models.e2e_asr_bat import BATModel

			from funasr.models.e2e_sa_asr import SAASRModel
			from funasr.models.e2e_asr_paraformer import Paraformer, ParaformerOnline, 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, ConformerChunkEncoder
			from funasr.models.encoder.data2vec_encoder import Data2VecEncoder
			from funasr.models.encoder.mfcca_encoder import MFCCAEncoder
			from funasr.models.encoder.resnet34_encoder import ResNet34Diar
			from funasr.models.encoder.rnn_encoder import RNNEncoder
			from funasr.models.encoder.sanm_encoder import SANMEncoder, SANMEncoderChunkOpt
			from funasr.models.encoder.branchformer_encoder import BranchformerEncoder
			from funasr.models.encoder.e_branchformer_encoder import EBranchformerEncoder
			from funasr.models.encoder.transformer_encoder import TransformerEncoder
			from funasr.models.encoder.rwkv_encoder import RWKVEncoder
			from funasr.models.frontend.default import DefaultFrontend
			from funasr.models.frontend.default import MultiChannelFrontend
			from funasr.models.frontend.fused import FusedFrontends
			@@ -40,7 +50,7 @@
			from funasr.models.frontend.wav_frontend import WavFrontend
			from funasr.models.frontend.windowing import SlidingWindow
			from funasr.models.joint_net.joint_network import JointNetwork
			from funasr.models.predictor.cif import CifPredictor, CifPredictorV2, CifPredictorV3
			from funasr.models.predictor.cif import CifPredictor, CifPredictorV2, CifPredictorV3, BATPredictor
			from funasr.models.specaug.specaug import SpecAug
			from funasr.models.specaug.specaug import SpecAugLFR
			from funasr.modules.subsampling import Conv1dSubsampling
			@@ -87,10 +97,13 @@
			paraformer_bert=ParaformerBert,
			bicif_paraformer=BiCifParaformer,
			contextual_paraformer=ContextualParaformer,
			neatcontextual_paraformer=NeatContextualParaformer,
			mfcca=MFCCA,
			timestamp_prediction=TimestampPredictor,
			rnnt=TransducerModel,
			rnnt_unified=UnifiedTransducerModel,
			sa_asr=SAASRModel,
			bat=BATModel,
			),
			default="asr",
			)
			@@ -103,10 +116,34 @@
			sanm=SANMEncoder,
			sanm_chunk_opt=SANMEncoderChunkOpt,
			data2vec_encoder=Data2VecEncoder,
			branchformer=BranchformerEncoder,
			e_branchformer=EBranchformerEncoder,
			mfcca_enc=MFCCAEncoder,
			chunk_conformer=ConformerChunkEncoder,
			rwkv=RWKVEncoder,
			),
			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,
			),
			default="rnn",
			)

			spk_encoder_choices = ClassChoices(
			"spk_encoder",
			classes=dict(
			resnet34_diar=ResNet34Diar,
			),
			default="resnet34_diar",
			)
			encoder_choices2 = ClassChoices(
			"encoder2",
			@@ -132,6 +169,7 @@
			paraformer_decoder_sanm=ParaformerSANMDecoder,
			paraformer_decoder_san=ParaformerDecoderSAN,
			contextual_paraformer_decoder=ContextualParaformerDecoder,
			sa_decoder=SAAsrTransformerDecoder,
			),
			default="rnn",
			)
			@@ -157,6 +195,7 @@
			ctc_predictor=None,
			cif_predictor_v2=CifPredictorV2,
			cif_predictor_v3=CifPredictorV3,
			bat_predictor=BATPredictor,
			),
			default="cif_predictor",
			optional=True,
			@@ -223,6 +262,10 @@
			rnnt_decoder_choices,
			# --joint_network and --joint_network_conf
			joint_network_choices,
			# --asr_encoder and --asr_encoder_conf
			asr_encoder_choices,
			# --spk_encoder and --spk_encoder_conf
			spk_encoder_choices,
			]


			@@ -245,7 +288,7 @@
			# frontend
			if hasattr(args, "input_size") and args.input_size is None:
			frontend_class = frontend_choices.get_class(args.frontend)
			if args.frontend == 'wav_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)
			@@ -267,7 +310,7 @@
			if args.normalize is not None:
			normalize_class = normalize_choices.get_class(args.normalize)
			if args.model == "mfcca":
			normalize = normalize_class(stats_file=args.cmvn_file,**args.normalize_conf)
			normalize = normalize_class(stats_file=args.cmvn_file, **args.normalize_conf)
			else:
			normalize = normalize_class(**args.normalize_conf)
			else:
			@@ -278,12 +321,15 @@
			encoder = encoder_class(input_size=input_size, **args.encoder_conf)

			# decoder
			decoder_class = decoder_choices.get_class(args.decoder)
			decoder = decoder_class(
			vocab_size=vocab_size,
			encoder_output_size=encoder.output_size(),
			**args.decoder_conf,
			)
			if hasattr(args, "decoder") and args.decoder is not None:
			decoder_class = decoder_choices.get_class(args.decoder)
			decoder = decoder_class(
			vocab_size=vocab_size,
			encoder_output_size=encoder.output_size(),
			**args.decoder_conf,
			)
			else:
			decoder = None

			# ctc
			ctc = CTC(
			@@ -373,10 +419,15 @@
			**args.model_conf,
			)
			elif args.model == "timestamp_prediction":
			# predictor
			predictor_class = predictor_choices.get_class(args.predictor)
			predictor = predictor_class(**args.predictor_conf)

			model_class = model_choices.get_class(args.model)
			model = model_class(
			frontend=frontend,
			encoder=encoder,
			predictor=predictor,
			token_list=token_list,
			**args.model_conf,
			)
			@@ -423,6 +474,80 @@
			joint_network=joint_network,
			**args.model_conf,
			)
			elif args.model == "bat":
			# 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,
			)

			predictor_class = predictor_choices.get_class(args.predictor)
			predictor = predictor_class(**args.predictor_conf)

			model_class = model_choices.get_class(args.model)
			# 7. Build model
			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,
			predictor=predictor,
			**args.model_conf,
			)
			elif args.model == "sa_asr":
			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)
			decoder = decoder_class(
			vocab_size=vocab_size,
			encoder_output_size=asr_encoder.output_size(),
			**args.decoder_conf,
			)
			ctc = CTC(
			odim=vocab_size, encoder_output_size=asr_encoder.output_size(), **args.ctc_conf
			)

			model_class = model_choices.get_class(args.model)
			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,
			)

			else:
			raise NotImplementedError("Not supported model: {}".format(args.model))