import os import torch import torch.nn as nn from funasr.export.utils.torch_function import MakePadMask from funasr.export.utils.torch_function import sequence_mask from funasr.modules.attention import MultiHeadedAttentionSANMDecoder from funasr.export.models.modules.multihead_att import MultiHeadedAttentionSANMDecoder as MultiHeadedAttentionSANMDecoder_export from funasr.modules.attention import MultiHeadedAttentionCrossAtt from funasr.export.models.modules.multihead_att import MultiHeadedAttentionCrossAtt as MultiHeadedAttentionCrossAtt_export from funasr.modules.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM from funasr.export.models.modules.feedforward import PositionwiseFeedForwardDecoderSANM as PositionwiseFeedForwardDecoderSANM_export from funasr.export.models.modules.decoder_layer import DecoderLayerSANM as DecoderLayerSANM_export class ContextualSANMDecoder(nn.Module): def __init__(self, model, max_seq_len=512, model_name='decoder', onnx: bool = True,): super().__init__() # self.embed = model.embed #Embedding(model.embed, max_seq_len) self.model = model if onnx: self.make_pad_mask = MakePadMask(max_seq_len, flip=False) else: self.make_pad_mask = sequence_mask(max_seq_len, flip=False) for i, d in enumerate(self.model.decoders): if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM): d.feed_forward = PositionwiseFeedForwardDecoderSANM_export(d.feed_forward) if isinstance(d.self_attn, MultiHeadedAttentionSANMDecoder): d.self_attn = MultiHeadedAttentionSANMDecoder_export(d.self_attn) if isinstance(d.src_attn, MultiHeadedAttentionCrossAtt): d.src_attn = MultiHeadedAttentionCrossAtt_export(d.src_attn) self.model.decoders[i] = DecoderLayerSANM_export(d) if self.model.decoders2 is not None: for i, d in enumerate(self.model.decoders2): if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM): d.feed_forward = PositionwiseFeedForwardDecoderSANM_export(d.feed_forward) if isinstance(d.self_attn, MultiHeadedAttentionSANMDecoder): d.self_attn = MultiHeadedAttentionSANMDecoder_export(d.self_attn) self.model.decoders2[i] = DecoderLayerSANM_export(d) for i, d in enumerate(self.model.decoders3): if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM): d.feed_forward = PositionwiseFeedForwardDecoderSANM_export(d.feed_forward) self.model.decoders3[i] = DecoderLayerSANM_export(d) self.output_layer = model.output_layer self.after_norm = model.after_norm self.model_name = model_name # bias decoder if isinstance(self.model.bias_decoder.src_attn, MultiHeadedAttentionCrossAtt): self.model.bias_decoder.src_attn = MultiHeadedAttentionCrossAtt_export(self.model.bias_decoder.src_attn) self.bias_decoder = self.model.bias_decoder # last decoder if isinstance(self.model.last_decoder.src_attn, MultiHeadedAttentionCrossAtt): self.model.last_decoder.src_attn = MultiHeadedAttentionCrossAtt_export(self.model.last_decoder.src_attn) if isinstance(self.model.last_decoder.self_attn, MultiHeadedAttentionSANMDecoder): self.model.last_decoder.self_attn = MultiHeadedAttentionSANMDecoder_export(self.model.last_decoder.self_attn) if isinstance(self.model.last_decoder.feed_forward, PositionwiseFeedForwardDecoderSANM): self.model.last_decoder.feed_forward = PositionwiseFeedForwardDecoderSANM_export(self.model.last_decoder.feed_forward) self.last_decoder = self.model.last_decoder self.bias_output = self.model.bias_output self.dropout = self.model.dropout def prepare_mask(self, mask): mask_3d_btd = mask[:, :, None] if len(mask.shape) == 2: mask_4d_bhlt = 1 - mask[:, None, None, :] elif len(mask.shape) == 3: mask_4d_bhlt = 1 - mask[:, None, :] mask_4d_bhlt = mask_4d_bhlt * -10000.0 return mask_3d_btd, mask_4d_bhlt def forward( self, hs_pad: torch.Tensor, hlens: torch.Tensor, ys_in_pad: torch.Tensor, ys_in_lens: torch.Tensor, bias_embed: torch.Tensor, ): tgt = ys_in_pad tgt_mask = self.make_pad_mask(ys_in_lens) tgt_mask, _ = self.prepare_mask(tgt_mask) # tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] memory = hs_pad memory_mask = self.make_pad_mask(hlens) _, memory_mask = self.prepare_mask(memory_mask) # memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] x = tgt x, tgt_mask, memory, memory_mask, _ = self.model.decoders( x, tgt_mask, memory, memory_mask ) _, _, x_self_attn, x_src_attn = self.last_decoder( x, tgt_mask, memory, memory_mask ) # contextual paraformer related contextual_length = torch.Tensor([bias_embed.shape[1]]).int().repeat(hs_pad.shape[0]) # contextual_mask = myutils.sequence_mask(contextual_length, device=memory.device)[:, None, :] contextual_mask = self.make_pad_mask(contextual_length) contextual_mask, _ = self.prepare_mask(contextual_mask) # import pdb; pdb.set_trace() contextual_mask = contextual_mask.transpose(2, 1).unsqueeze(1) cx, tgt_mask, _, _, _ = self.bias_decoder(x_self_attn, tgt_mask, bias_embed, memory_mask=contextual_mask) if self.bias_output is not None: x = torch.cat([x_src_attn, cx], dim=2) x = self.bias_output(x.transpose(1, 2)).transpose(1, 2) # 2D -> D x = x_self_attn + self.dropout(x) if self.model.decoders2 is not None: x, tgt_mask, memory, memory_mask, _ = self.model.decoders2( x, tgt_mask, memory, memory_mask ) x, tgt_mask, memory, memory_mask, _ = self.model.decoders3( x, tgt_mask, memory, memory_mask ) x = self.after_norm(x) x = self.output_layer(x) return x, ys_in_lens def get_dummy_inputs(self, enc_size): tgt = torch.LongTensor([0]).unsqueeze(0) memory = torch.randn(1, 100, enc_size) pre_acoustic_embeds = torch.randn(1, 1, enc_size) cache_num = len(self.model.decoders) + len(self.model.decoders2) cache = [ torch.zeros((1, self.model.decoders[0].size, self.model.decoders[0].self_attn.kernel_size)) for _ in range(cache_num) ] return (tgt, memory, pre_acoustic_embeds, cache) def is_optimizable(self): return True def get_input_names(self): cache_num = len(self.model.decoders) + len(self.model.decoders2) return ['tgt', 'memory', 'pre_acoustic_embeds'] \ + ['cache_%d' % i for i in range(cache_num)] def get_output_names(self): cache_num = len(self.model.decoders) + len(self.model.decoders2) return ['y'] \ + ['out_cache_%d' % i for i in range(cache_num)] def get_dynamic_axes(self): ret = { 'tgt': { 0: 'tgt_batch', 1: 'tgt_length' }, 'memory': { 0: 'memory_batch', 1: 'memory_length' }, 'pre_acoustic_embeds': { 0: 'acoustic_embeds_batch', 1: 'acoustic_embeds_length', } } cache_num = len(self.model.decoders) + len(self.model.decoders2) ret.update({ 'cache_%d' % d: { 0: 'cache_%d_batch' % d, 2: 'cache_%d_length' % d } for d in range(cache_num) }) return ret def get_model_config(self, path): return { "dec_type": "XformerDecoder", "model_path": os.path.join(path, f'{self.model_name}.onnx'), "n_layers": len(self.model.decoders) + len(self.model.decoders2), "odim": self.model.decoders[0].size }