mirror of
https://github.com/modelscope/FunASR
synced 2025-09-15 14:48:36 +08:00
add
This commit is contained in:
mengzhe.cmz
parent
1ad439f96b
commit
e21a6ed2d8
@ -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):
|
||||
|
||||
@ -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'
|
||||
},
|
||||
}
|
||||
@ -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,
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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",
|
||||
|
||||
Loading…
Reference in New Issue
Block a user