From 4ace5a95b052d338947fc88809a440ccd55cf6b4 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 16 十一月 2023 16:39:52 +0800
Subject: [PATCH] funasr pages

---
 funasr/export/models/encoder/sanm_encoder.py |   60 +++++++++++++++++++++++++++++++++---------------------------
 1 files changed, 33 insertions(+), 27 deletions(-)

diff --git a/funasr/export/models/encoder/sanm_encoder.py b/funasr/export/models/encoder/sanm_encoder.py
index 44a48ff..d1b4b1e 100644
--- a/funasr/export/models/encoder/sanm_encoder.py
+++ b/funasr/export/models/encoder/sanm_encoder.py
@@ -8,6 +8,7 @@
 from funasr.export.models.modules.encoder_layer import EncoderLayerSANM as EncoderLayerSANM_export
 from funasr.modules.positionwise_feed_forward import PositionwiseFeedForward
 from funasr.export.models.modules.feedforward import PositionwiseFeedForward as PositionwiseFeedForward_export
+from funasr.modules.embedding import StreamSinusoidalPositionEncoder
 
 
 class SANMEncoder(nn.Module):
@@ -21,6 +22,8 @@
     ):
         super().__init__()
         self.embed = model.embed
+        if isinstance(self.embed, StreamSinusoidalPositionEncoder):
+            self.embed = None
         self.model = model
         self.feats_dim = feats_dim
         self._output_size = model._output_size
@@ -63,8 +66,10 @@
     def forward(self,
                 speech: torch.Tensor,
                 speech_lengths: torch.Tensor,
+                online: bool = False 
                 ):
-        speech = speech * self._output_size ** 0.5
+        if not online:
+            speech = speech * self._output_size ** 0.5
         mask = self.make_pad_mask(speech_lengths)
         mask = self.prepare_mask(mask)
         if self.embed is None:
@@ -158,13 +163,14 @@
     def forward(self,
                 speech: torch.Tensor,
                 speech_lengths: torch.Tensor,
-                vad_mask: torch.Tensor,
+                vad_masks: torch.Tensor,
                 sub_masks: torch.Tensor,
                 ):
         speech = speech * self._output_size ** 0.5
         mask = self.make_pad_mask(speech_lengths)
+        vad_masks = self.prepare_mask(mask, vad_masks)
         mask = self.prepare_mask(mask, sub_masks)
-        vad_mask = self.prepare_mask(mask, vad_mask)
+        
         if self.embed is None:
             xs_pad = speech
         else:
@@ -176,7 +182,7 @@
         # encoder_outs = self.model.encoders(xs_pad, mask)
         for layer_idx, encoder_layer in enumerate(self.model.encoders):
             if layer_idx == len(self.model.encoders) - 1:
-                mask = vad_mask
+                mask = vad_masks
             encoder_outs = encoder_layer(xs_pad, mask)
             xs_pad, masks = encoder_outs[0], encoder_outs[1]
         
@@ -187,26 +193,26 @@
     def get_output_size(self):
         return self.model.encoders[0].size
     
-    def get_dummy_inputs(self):
-        feats = torch.randn(1, 100, self.feats_dim)
-        return (feats)
-    
-    def get_input_names(self):
-        return ['feats']
-    
-    def get_output_names(self):
-        return ['encoder_out', 'encoder_out_lens', 'predictor_weight']
-    
-    def get_dynamic_axes(self):
-        return {
-            'feats': {
-                1: 'feats_length'
-            },
-            'encoder_out': {
-                1: 'enc_out_length'
-            },
-            'predictor_weight': {
-                1: 'pre_out_length'
-            }
-            
-        }
+    # def get_dummy_inputs(self):
+    #     feats = torch.randn(1, 100, self.feats_dim)
+    #     return (feats)
+    #
+    # def get_input_names(self):
+    #     return ['feats']
+    #
+    # def get_output_names(self):
+    #     return ['encoder_out', 'encoder_out_lens', 'predictor_weight']
+    #
+    # def get_dynamic_axes(self):
+    #     return {
+    #         'feats': {
+    #             1: 'feats_length'
+    #         },
+    #         'encoder_out': {
+    #             1: 'enc_out_length'
+    #         },
+    #         'predictor_weight': {
+    #             1: 'pre_out_length'
+    #         }
+    #
+    #     }

--
Gitblit v1.9.1