From e21a6ed2d81e7cd07d50ec1b6f6127b3d638a27b Mon Sep 17 00:00:00 2001
From: mengzhe.cmz <mengzhe.cmz@alibaba-inc.com>
Date: 星期四, 13 四月 2023 13:38:22 +0800
Subject: [PATCH] add
---
/dev/null | 154 ---------------------------------------------------
funasr/models/target_delay_transformer.py | 6 +
funasr/runtime/python/onnxruntime/funasr_onnx/punc_bin.py | 10 +++
funasr/export/models/__init__.py | 4
funasr/models/vad_realtime_transformer.py | 6 +
5 files changed, 22 insertions(+), 158 deletions(-)
diff --git a/funasr/export/models/__init__.py b/funasr/export/models/__init__.py
index f81ff64..0e3a782 100644
--- a/funasr/export/models/__init__.py
+++ b/funasr/export/models/__init__.py
@@ -4,10 +4,10 @@
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):
diff --git a/funasr/export/models/target_delay_transformer.py b/funasr/export/models/target_delay_transformer.py
deleted file mode 100644
index 2780d82..0000000
--- a/funasr/export/models/target_delay_transformer.py
+++ /dev/null
@@ -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'
- },
- }
diff --git a/funasr/models/target_delay_transformer.py b/funasr/models/target_delay_transformer.py
index 84a2e6c..8cd4357 100644
--- a/funasr/models/target_delay_transformer.py
+++ b/funasr/models/target_delay_transformer.py
@@ -13,7 +13,11 @@
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,
diff --git a/funasr/models/vad_realtime_transformer.py b/funasr/models/vad_realtime_transformer.py
index 66f7fad..3810672 100644
--- a/funasr/models/vad_realtime_transformer.py
+++ b/funasr/models/vad_realtime_transformer.py
@@ -11,7 +11,11 @@
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,
diff --git a/funasr/runtime/python/onnxruntime/funasr_onnx/punc_bin.py b/funasr/runtime/python/onnxruntime/funasr_onnx/punc_bin.py
index 0eb764f..2f1b3b7 100644
--- a/funasr/runtime/python/onnxruntime/funasr_onnx/punc_bin.py
+++ b/funasr/runtime/python/onnxruntime/funasr_onnx/punc_bin.py
@@ -13,6 +13,11 @@
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_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",
--
Gitblit v1.9.1