From 9d01231fa69672fc7c9b4bf81ef466bb0189788c Mon Sep 17 00:00:00 2001
From: aky15 <ankeyu.aky@11.17.44.249>
Date: 星期三, 17 五月 2023 17:34:21 +0800
Subject: [PATCH] rnnt继承ASRTask
---
funasr/build_utils/build_asr_model.py | 86 ++++++++++++++++++++++++++++++++++++++++++-
1 files changed, 84 insertions(+), 2 deletions(-)
diff --git a/funasr/build_utils/build_asr_model.py b/funasr/build_utils/build_asr_model.py
index d8cbba5..718736b 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 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 @@
sanm_chunk_opt=SANMEncoderChunkOpt,
data2vec_encoder=Data2VecEncoder,
mfcca_enc=MFCCAEncoder,
+ chunk_conformer=ConformerChunkEncoder,
),
default="rnn",
)
@@ -171,6 +175,23 @@
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 @@
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 @@
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 @@
if args.init is not None:
initialize(model, args.init)
- return model
\ No newline at end of file
+ return model
--
Gitblit v1.9.1