From 52eb056c76f80e1bc8e77867de8922bb6a1c41c6 Mon Sep 17 00:00:00 2001
From: speech_asr <wangjiaming.wjm@alibaba-inc.com>
Date: 星期五, 21 四月 2023 01:29:44 +0800
Subject: [PATCH] update
---
funasr/build_utils/build_model.py | 3
funasr/build_utils/build_diar_model.py | 296 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 299 insertions(+), 0 deletions(-)
diff --git a/funasr/build_utils/build_diar_model.py b/funasr/build_utils/build_diar_model.py
new file mode 100644
index 0000000..6406404
--- /dev/null
+++ b/funasr/build_utils/build_diar_model.py
@@ -0,0 +1,296 @@
+import logging
+
+import torch
+
+from funasr.layers.global_mvn import GlobalMVN
+from funasr.layers.label_aggregation import LabelAggregate
+from funasr.layers.utterance_mvn import UtteranceMVN
+from funasr.models.e2e_diar_eend_ola import DiarEENDOLAModel
+from funasr.models.e2e_diar_sond import DiarSondModel
+from funasr.models.encoder.conformer_encoder import ConformerEncoder
+from funasr.models.encoder.data2vec_encoder import Data2VecEncoder
+from funasr.models.encoder.ecapa_tdnn_encoder import ECAPA_TDNN
+from funasr.models.encoder.opennmt_encoders.ci_scorers import DotScorer, CosScorer
+from funasr.models.encoder.opennmt_encoders.conv_encoder import ConvEncoder
+from funasr.models.encoder.opennmt_encoders.fsmn_encoder import FsmnEncoder
+from funasr.models.encoder.opennmt_encoders.self_attention_encoder import SelfAttentionEncoder
+from funasr.models.encoder.resnet34_encoder import ResNet34Diar, ResNet34SpL2RegDiar
+from funasr.models.encoder.rnn_encoder import RNNEncoder
+from funasr.models.encoder.sanm_encoder import SANMEncoder, SANMEncoderChunkOpt
+from funasr.models.encoder.transformer_encoder import TransformerEncoder
+from funasr.models.frontend.default import DefaultFrontend
+from funasr.models.frontend.fused import FusedFrontends
+from funasr.models.frontend.s3prl import S3prlFrontend
+from funasr.models.frontend.wav_frontend import WavFrontend
+from funasr.models.frontend.wav_frontend import WavFrontendMel23
+from funasr.models.frontend.windowing import SlidingWindow
+from funasr.models.specaug.specaug import SpecAug
+from funasr.models.specaug.specaug import SpecAugLFR
+from funasr.modules.eend_ola.encoder import EENDOLATransformerEncoder
+from funasr.modules.eend_ola.encoder_decoder_attractor import EncoderDecoderAttractor
+from funasr.torch_utils.initialize import initialize
+from funasr.train.class_choices import ClassChoices
+
+frontend_choices = ClassChoices(
+ name="frontend",
+ classes=dict(
+ default=DefaultFrontend,
+ sliding_window=SlidingWindow,
+ s3prl=S3prlFrontend,
+ fused=FusedFrontends,
+ wav_frontend=WavFrontend,
+ wav_frontend_mel23=WavFrontendMel23,
+ ),
+ default="default",
+)
+specaug_choices = ClassChoices(
+ name="specaug",
+ classes=dict(
+ specaug=SpecAug,
+ specaug_lfr=SpecAugLFR,
+ ),
+ default=None,
+ optional=True,
+)
+normalize_choices = ClassChoices(
+ "normalize",
+ classes=dict(
+ global_mvn=GlobalMVN,
+ utterance_mvn=UtteranceMVN,
+ ),
+ default=None,
+ optional=True,
+)
+label_aggregator_choices = ClassChoices(
+ "label_aggregator",
+ classes=dict(
+ label_aggregator=LabelAggregate
+ ),
+ default=None,
+ optional=True,
+)
+model_choices = ClassChoices(
+ "model",
+ classes=dict(
+ sond=DiarSondModel,
+ eend_ola=DiarEENDOLAModel,
+ ),
+ default="sond",
+)
+encoder_choices = ClassChoices(
+ "encoder",
+ classes=dict(
+ conformer=ConformerEncoder,
+ transformer=TransformerEncoder,
+ rnn=RNNEncoder,
+ sanm=SANMEncoder,
+ san=SelfAttentionEncoder,
+ fsmn=FsmnEncoder,
+ conv=ConvEncoder,
+ resnet34=ResNet34Diar,
+ resnet34_sp_l2reg=ResNet34SpL2RegDiar,
+ sanm_chunk_opt=SANMEncoderChunkOpt,
+ data2vec_encoder=Data2VecEncoder,
+ ecapa_tdnn=ECAPA_TDNN,
+ eend_ola_transformer=EENDOLATransformerEncoder,
+ ),
+ default="resnet34",
+)
+speaker_encoder_choices = ClassChoices(
+ "speaker_encoder",
+ classes=dict(
+ conformer=ConformerEncoder,
+ transformer=TransformerEncoder,
+ rnn=RNNEncoder,
+ sanm=SANMEncoder,
+ san=SelfAttentionEncoder,
+ fsmn=FsmnEncoder,
+ conv=ConvEncoder,
+ sanm_chunk_opt=SANMEncoderChunkOpt,
+ data2vec_encoder=Data2VecEncoder,
+ ),
+ default=None,
+ optional=True
+)
+cd_scorer_choices = ClassChoices(
+ "cd_scorer",
+ classes=dict(
+ san=SelfAttentionEncoder,
+ ),
+ default=None,
+ optional=True,
+)
+ci_scorer_choices = ClassChoices(
+ "ci_scorer",
+ classes=dict(
+ dot=DotScorer,
+ cosine=CosScorer,
+ conv=ConvEncoder,
+ ),
+ type_check=torch.nn.Module,
+ default=None,
+ optional=True,
+)
+# decoder is used for output (e.g. post_net in SOND)
+decoder_choices = ClassChoices(
+ "decoder",
+ classes=dict(
+ rnn=RNNEncoder,
+ fsmn=FsmnEncoder,
+ ),
+ type_check=torch.nn.Module,
+ default="fsmn",
+)
+# encoder_decoder_attractor is used for EEND-OLA
+encoder_decoder_attractor_choices = ClassChoices(
+ "encoder_decoder_attractor",
+ classes=dict(
+ eda=EncoderDecoderAttractor,
+ ),
+ type_check=torch.nn.Module,
+ default="eda",
+)
+class_choices_list = [
+ # --frontend and --frontend_conf
+ frontend_choices,
+ # --specaug and --specaug_conf
+ specaug_choices,
+ # --normalize and --normalize_conf
+ normalize_choices,
+ # --label_aggregator and --label_aggregator_conf
+ label_aggregator_choices,
+ # --model and --model_conf
+ model_choices,
+ # --encoder and --encoder_conf
+ encoder_choices,
+ # --speaker_encoder and --speaker_encoder_conf
+ speaker_encoder_choices,
+ # --cd_scorer and cd_scorer_conf
+ cd_scorer_choices,
+ # --ci_scorer and ci_scorer_conf
+ ci_scorer_choices,
+ # --decoder and --decoder_conf
+ decoder_choices,
+ # --eda and --eda_conf
+ encoder_decoder_attractor_choices,
+]
+
+
+def build_diar_model(args):
+ # token_list
+ if args.token_list is not None:
+ with open(args.token_list) as f:
+ token_list = [line.rstrip() for line in f]
+ args.token_list = list(token_list)
+ vocab_size = len(token_list)
+ logging.info(f"Vocabulary size: {vocab_size}")
+ else:
+ vocab_size = None
+
+ # frontend
+ if args.input_size is None:
+ frontend_class = frontend_choices.get_class(args.frontend)
+ if args.frontend == 'wav_frontend':
+ frontend = frontend_class(cmvn_file=args.cmvn_file, **args.frontend_conf)
+ else:
+ frontend = frontend_class(**args.frontend_conf)
+ input_size = frontend.output_size()
+ else:
+ args.frontend = None
+ args.frontend_conf = {}
+ frontend = None
+ input_size = args.input_size
+
+ # encoder
+ encoder_class = encoder_choices.get_class(args.encoder)
+ encoder = encoder_class(input_size=input_size, **args.encoder_conf)
+
+ if args.model_name == "sond":
+ # 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
+
+ # 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
+
+ # speaker encoder
+ if getattr(args, "speaker_encoder", None) is not None:
+ speaker_encoder_class = speaker_encoder_choices.get_class(args.speaker_encoder)
+ speaker_encoder = speaker_encoder_class(**args.speaker_encoder_conf)
+ else:
+ speaker_encoder = None
+
+ # ci scorer
+ if getattr(args, "ci_scorer", None) is not None:
+ ci_scorer_class = ci_scorer_choices.get_class(args.ci_scorer)
+ ci_scorer = ci_scorer_class(**args.ci_scorer_conf)
+ else:
+ ci_scorer = None
+
+ # cd scorer
+ if getattr(args, "cd_scorer", None) is not None:
+ cd_scorer_class = cd_scorer_choices.get_class(args.cd_scorer)
+ cd_scorer = cd_scorer_class(**args.cd_scorer_conf)
+ else:
+ cd_scorer = None
+
+ # 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,
+ )
+
+ # logger aggregator
+ if getattr(args, "label_aggregator", None) is not None:
+ label_aggregator_class = label_aggregator_choices.get_class(args.label_aggregator)
+ label_aggregator = label_aggregator_class(**args.label_aggregator_conf)
+ else:
+ label_aggregator = None
+
+ model_class = model_choices.get_class(args.model)
+ model = model_class(
+ vocab_size=vocab_size,
+ frontend=frontend,
+ specaug=specaug,
+ normalize=normalize,
+ label_aggregator=label_aggregator,
+ encoder=encoder,
+ speaker_encoder=speaker_encoder,
+ ci_scorer=ci_scorer,
+ cd_scorer=cd_scorer,
+ decoder=decoder,
+ token_list=token_list,
+ **args.model_conf,
+ )
+
+ elif args.model_name == "eend_ola":
+ # encoder-decoder attractor
+ encoder_decoder_attractor_class = encoder_decoder_attractor_choices.get_class(args.encoder_decoder_attractor)
+ encoder_decoder_attractor = encoder_decoder_attractor_class(**args.encoder_decoder_attractor_conf)
+
+ # 9. Build model
+ model_class = model_choices.get_class(args.model)
+ model = model_class(
+ frontend=frontend,
+ encoder=encoder,
+ encoder_decoder_attractor=encoder_decoder_attractor,
+ **args.model_conf,
+ )
+
+ else:
+ raise NotImplementedError("Not supported model: {}".format(args.model))
+
+ # 10. Initialize
+ if args.init is not None:
+ initialize(model, args.init)
+
+ return model
diff --git a/funasr/build_utils/build_model.py b/funasr/build_utils/build_model.py
index 6029fae..13a6faa 100644
--- a/funasr/build_utils/build_model.py
+++ b/funasr/build_utils/build_model.py
@@ -3,6 +3,7 @@
from funasr.build_utils.build_pretrain_model import build_pretrain_model
from funasr.build_utils.build_punc_model import build_punc_model
from funasr.build_utils.build_vad_model import build_vad_model
+from funasr.build_utils.build_diar_model import build_diar_model
def build_model(args):
@@ -16,6 +17,8 @@
model = build_punc_model(args)
elif args.task_name == "vad":
model = build_vad_model(args)
+ elif args.task_name == "diar":
+ model = build_diar_model(args)
else:
raise NotImplementedError("Not supported task: {}".format(args.task_name))
--
Gitblit v1.9.1