add

funasr/export/models/__init__.py

@@ -4,10 +4,10 @@ from funasr.export.models.e2e_asr_paraformer import BiCifParaformer as BiCifPara
 from funasr.models.e2e_vad import E2EVadModel
 from funasr.export.models.e2e_vad import E2EVadModel as E2EVadModel_export
 from funasr.models.target_delay_transformer import TargetDelayTransformer
-from funasr.export.models.target_delay_transformer import CT_Transformer as CT_Transformer_export
+from funasr.export.models.CT_Transformer import CT_Transformer as CT_Transformer_export
 from funasr.train.abs_model import PunctuationModel
 from funasr.models.vad_realtime_transformer import VadRealtimeTransformer
-from funasr.export.models.target_delay_transformer import CT_Transformer_VadRealtime as CT_Transformer_VadRealtime_export
+from funasr.export.models.CT_Transformer import CT_Transformer_VadRealtime as CT_Transformer_VadRealtime_export
 
 def get_model(model, export_config=None):
     if isinstance(model, BiCifParaformer):

funasr/export/models/target_delay_transformer.py

@@ -1,154 +0,0 @@
-from typing import Tuple
-
-import torch
-import torch.nn as nn
-
-from funasr.models.encoder.sanm_encoder import SANMEncoder
-from funasr.export.models.encoder.sanm_encoder import SANMEncoder as SANMEncoder_export
-from funasr.models.encoder.sanm_encoder import SANMVadEncoder
-from funasr.export.models.encoder.sanm_encoder import SANMVadEncoder as SANMVadEncoder_export
-
-class CT_Transformer(nn.Module):
-
-    def __init__(
-            self,
-            model,
-            max_seq_len=512,
-            model_name='punc_model',
-            **kwargs,
-    ):
-        super().__init__()
-        onnx = False
-        if "onnx" in kwargs:
-            onnx = kwargs["onnx"]
-        self.embed = model.embed
-        self.decoder = model.decoder
-        # self.model = model
-        self.feats_dim = self.embed.embedding_dim
-        self.num_embeddings = self.embed.num_embeddings
-        self.model_name = model_name
-
-        if isinstance(model.encoder, SANMEncoder):
-            self.encoder = SANMEncoder_export(model.encoder, onnx=onnx)
-        else:
-            assert False, "Only support samn encode."
-
-    def forward(self, inputs: torch.Tensor, text_lengths: torch.Tensor) -> Tuple[torch.Tensor, None]:
-        """Compute loss value from buffer sequences.
-
-        Args:
-            input (torch.Tensor): Input ids. (batch, len)
-            hidden (torch.Tensor): Target ids. (batch, len)
-
-        """
-        x = self.embed(inputs)
-        # mask = self._target_mask(input)
-        h, _ = self.encoder(x, text_lengths)
-        y = self.decoder(h)
-        return y
-
-    def get_dummy_inputs(self):
-        length = 120
-        text_indexes = torch.randint(0, self.embed.num_embeddings, (2, length))
-        text_lengths = torch.tensor([length-20, length], dtype=torch.int32)
-        return (text_indexes, text_lengths)
-
-    def get_input_names(self):
-        return ['inputs', 'text_lengths']
-
-    def get_output_names(self):
-        return ['logits']
-
-    def get_dynamic_axes(self):
-        return {
-            'inputs': {
-                0: 'batch_size',
-                1: 'feats_length'
-            },
-            'text_lengths': {
-                0: 'batch_size',
-            },
-            'logits': {
-                0: 'batch_size',
-                1: 'logits_length'
-            },
-        }
-
-
-class CT_Transformer_VadRealtime(nn.Module):
-
-    def __init__(
-        self,
-        model,
-        max_seq_len=512,
-        model_name='punc_model',
-        **kwargs,
-    ):
-        super().__init__()
-        onnx = False
-        if "onnx" in kwargs:
-            onnx = kwargs["onnx"]
-
-        self.embed = model.embed
-        if isinstance(model.encoder, SANMVadEncoder):
-            self.encoder = SANMVadEncoder_export(model.encoder, onnx=onnx)
-        else:
-            assert False, "Only support samn encode."
-        self.decoder = model.decoder
-        self.model_name = model_name
-
-
-
-    def forward(self, inputs: torch.Tensor,
-                text_lengths: torch.Tensor,
-                vad_indexes: torch.Tensor,
-                sub_masks: torch.Tensor,
-                ) -> Tuple[torch.Tensor, None]:
-        """Compute loss value from buffer sequences.
-
-        Args:
-            input (torch.Tensor): Input ids. (batch, len)
-            hidden (torch.Tensor): Target ids. (batch, len)
-
-        """
-        x = self.embed(inputs)
-        # mask = self._target_mask(input)
-        h, _ = self.encoder(x, text_lengths, vad_indexes, sub_masks)
-        y = self.decoder(h)
-        return y
-
-    def with_vad(self):
-        return True
-
-    def get_dummy_inputs(self):
-        length = 120
-        text_indexes = torch.randint(0, self.embed.num_embeddings, (1, length))
-        text_lengths = torch.tensor([length], dtype=torch.int32)
-        vad_mask = torch.ones(length, length, dtype=torch.float32)[None, None, :, :]
-        sub_masks = torch.ones(length, length, dtype=torch.float32)
-        sub_masks = torch.tril(sub_masks).type(torch.float32)
-        return (text_indexes, text_lengths, vad_mask, sub_masks[None, None, :, :])
-
-    def get_input_names(self):
-        return ['inputs', 'text_lengths', 'vad_masks', 'sub_masks']
-
-    def get_output_names(self):
-        return ['logits']
-
-    def get_dynamic_axes(self):
-        return {
-            'inputs': {
-                1: 'feats_length'
-            },
-            'vad_masks': {
-                2: 'feats_length1',
-                3: 'feats_length2'
-            },
-            'sub_masks': {
-                2: 'feats_length1',
-                3: 'feats_length2'
-            },
-            'logits': {
-                1: 'logits_length'
-            },
-        }

funasr/models/target_delay_transformer.py

@@ -13,7 +13,11 @@ from funasr.train.abs_model import AbsPunctuation
 
 
 class TargetDelayTransformer(AbsPunctuation):
-
+    """
+    Author: Speech Lab, Alibaba Group, China
+    CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection
+    https://arxiv.org/pdf/2003.01309.pdf
+    """
     def __init__(
         self,
         vocab_size: int,

funasr/models/vad_realtime_transformer.py

@@ -11,7 +11,11 @@ from funasr.train.abs_model import AbsPunctuation
 
 
 class VadRealtimeTransformer(AbsPunctuation):
-
+    """
+    Author: Speech Lab, Alibaba Group, China
+    CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection
+    https://arxiv.org/pdf/2003.01309.pdf
+    """
     def __init__(
         self,
         vocab_size: int,

funasr/runtime/python/onnxruntime/funasr_onnx/punc_bin.py

@@ -13,6 +13,11 @@ logging = get_logger()
 
 
 class CT_Transformer():
+    """
+    Author: Speech Lab, Alibaba Group, China
+    CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection
+    https://arxiv.org/pdf/2003.01309.pdf
+    """
     def __init__(self, model_dir: Union[str, Path] = None,
                  batch_size: int = 1,
                  device_id: Union[str, int] = "-1",
@@ -119,6 +124,11 @@ class CT_Transformer():
 
 
 class CT_Transformer_VadRealtime(CT_Transformer):
+    """
+    Author: Speech Lab, Alibaba Group, China
+    CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection
+    https://arxiv.org/pdf/2003.01309.pdf
+    """
     def __init__(self, model_dir: Union[str, Path] = None,
                  batch_size: int = 1,
                  device_id: Union[str, int] = "-1",