first commit for takway.ai
This commit is contained in:
IrvingGao committed
commit
d7a6de771c
3256 files changed
+795173
No files matched your search
@@ -0,0 +1,273 @@
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
import triton_python_backend_utils as pb_utils
|
||||
import numpy as np
|
||||
from torch.utils.dlpack import from_dlpack
|
||||
import json
|
||||
import yaml
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
|
||||
|
||||
class LimitedDict(OrderedDict):
|
||||
def __init__(self, max_length):
|
||||
super().__init__()
|
||||
self.max_length = max_length
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if len(self) >= self.max_length:
|
||||
self.popitem(last=False)
|
||||
super().__setitem__(key, value)
|
||||
|
||||
|
||||
class CIFSearch:
|
||||
"""CIFSearch: https://github.com/alibaba-damo-academy/FunASR/blob/main/runtime/python/onnxruntime/funasr_onnx
|
||||
/paraformer_online_bin.py"""
|
||||
|
||||
def __init__(self):
|
||||
self.cache = {
|
||||
"cif_hidden": np.zeros((1, 1, 512)).astype(np.float32),
|
||||
"cif_alphas": np.zeros((1, 1)).astype(np.float32),
|
||||
"last_chunk": False,
|
||||
}
|
||||
self.chunk_size = [5, 10, 5]
|
||||
self.tail_threshold = 0.45
|
||||
self.cif_threshold = 1.0
|
||||
|
||||
def infer(self, hidden, alphas):
|
||||
batch_size, len_time, hidden_size = hidden.shape
|
||||
token_length = []
|
||||
list_fires = []
|
||||
list_frames = []
|
||||
cache_alphas = []
|
||||
cache_hiddens = []
|
||||
alphas[:, : self.chunk_size[0]] = 0.0
|
||||
alphas[:, sum(self.chunk_size[:2]) :] = 0.0
|
||||
|
||||
if self.cache is not None and "cif_alphas" in self.cache and "cif_hidden" in self.cache:
|
||||
hidden = np.concatenate((self.cache["cif_hidden"], hidden), axis=1)
|
||||
alphas = np.concatenate((self.cache["cif_alphas"], alphas), axis=1)
|
||||
if self.cache is not None and "last_chunk" in self.cache and self.cache["last_chunk"]:
|
||||
tail_hidden = np.zeros((batch_size, 1, hidden_size)).astype(np.float32)
|
||||
tail_alphas = np.array([[self.tail_threshold]]).astype(np.float32)
|
||||
tail_alphas = np.tile(tail_alphas, (batch_size, 1))
|
||||
hidden = np.concatenate((hidden, tail_hidden), axis=1)
|
||||
alphas = np.concatenate((alphas, tail_alphas), axis=1)
|
||||
|
||||
len_time = alphas.shape[1]
|
||||
for b in range(batch_size):
|
||||
integrate = 0.0
|
||||
frames = np.zeros(hidden_size).astype(np.float32)
|
||||
list_frame = []
|
||||
list_fire = []
|
||||
for t in range(len_time):
|
||||
alpha = alphas[b][t]
|
||||
if alpha + integrate < self.cif_threshold:
|
||||
integrate += alpha
|
||||
list_fire.append(integrate)
|
||||
frames += alpha * hidden[b][t]
|
||||
else:
|
||||
frames += (self.cif_threshold - integrate) * hidden[b][t]
|
||||
list_frame.append(frames)
|
||||
integrate += alpha
|
||||
list_fire.append(integrate)
|
||||
integrate -= self.cif_threshold
|
||||
frames = integrate * hidden[b][t]
|
||||
|
||||
cache_alphas.append(integrate)
|
||||
if integrate > 0.0:
|
||||
cache_hiddens.append(frames / integrate)
|
||||
else:
|
||||
cache_hiddens.append(frames)
|
||||
|
||||
token_length.append(len(list_frame))
|
||||
list_fires.append(list_fire)
|
||||
list_frames.append(list_frame)
|
||||
|
||||
max_token_len = max(token_length)
|
||||
list_ls = []
|
||||
for b in range(batch_size):
|
||||
pad_frames = np.zeros((max_token_len - token_length[b], hidden_size)).astype(np.float32)
|
||||
if token_length[b] == 0:
|
||||
list_ls.append(pad_frames)
|
||||
else:
|
||||
list_ls.append(np.concatenate((list_frames[b], pad_frames), axis=0))
|
||||
|
||||
self.cache["cif_alphas"] = np.stack(cache_alphas, axis=0)
|
||||
self.cache["cif_alphas"] = np.expand_dims(self.cache["cif_alphas"], axis=0)
|
||||
self.cache["cif_hidden"] = np.stack(cache_hiddens, axis=0)
|
||||
self.cache["cif_hidden"] = np.expand_dims(self.cache["cif_hidden"], axis=0)
|
||||
|
||||
return np.stack(list_ls, axis=0).astype(np.float32), np.stack(token_length, axis=0).astype(
|
||||
np.int32
|
||||
)
|
||||
|
||||
|
||||
class TritonPythonModel:
|
||||
"""Your Python model must use the same class name. Every Python model
|
||||
that is created must have "TritonPythonModel" as the class name.
|
||||
"""
|
||||
|
||||
def initialize(self, args):
|
||||
"""`initialize` is called only once when the model is being loaded.
|
||||
Implementing `initialize` function is optional. This function allows
|
||||
the model to initialize any state associated with this model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
args : dict
|
||||
Both keys and values are strings. The dictionary keys and values are:
|
||||
* model_config: A JSON string containing the model configuration
|
||||
* model_instance_kind: A string containing model instance kind
|
||||
* model_instance_device_id: A string containing model instance device ID
|
||||
* model_repository: Model repository path
|
||||
* model_version: Model version
|
||||
* model_name: Model name
|
||||
"""
|
||||
self.model_config = model_config = json.loads(args["model_config"])
|
||||
self.max_batch_size = max(model_config["max_batch_size"], 1)
|
||||
|
||||
# # Get OUTPUT0 configuration
|
||||
output0_config = pb_utils.get_output_config_by_name(model_config, "transcripts")
|
||||
# # Convert Triton types to numpy types
|
||||
self.out0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"])
|
||||
|
||||
self.init_vocab(self.model_config["parameters"])
|
||||
|
||||
self.cif_search_cache = LimitedDict(1024)
|
||||
self.start = LimitedDict(1024)
|
||||
|
||||
def init_vocab(self, parameters):
|
||||
for li in parameters.items():
|
||||
key, value = li
|
||||
value = value["string_value"]
|
||||
if key == "vocabulary":
|
||||
self.vocab_dict = self.load_vocab(value)
|
||||
|
||||
def load_vocab(self, vocab_file):
|
||||
with open(str(vocab_file), "rb") as f:
|
||||
config = yaml.load(f, Loader=yaml.Loader)
|
||||
return config["token_list"]
|
||||
|
||||
async def execute(self, requests):
|
||||
"""`execute` must be implemented in every Python model. `execute`
|
||||
function receives a list of pb_utils.InferenceRequest as the only
|
||||
argument. This function is called when an inference is requested
|
||||
for this model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
requests : list
|
||||
A list of pb_utils.InferenceRequest
|
||||
|
||||
Returns
|
||||
-------
|
||||
list
|
||||
A list of pb_utils.InferenceResponse. The length of this list must
|
||||
be the same as `requests`
|
||||
"""
|
||||
# Every Python backend must iterate through list of requests and create
|
||||
# an instance of pb_utils.InferenceResponse class for each of them. You
|
||||
# should avoid storing any of the input Tensors in the class attributes
|
||||
# as they will be overridden in subsequent inference requests. You can
|
||||
# make a copy of the underlying NumPy array and store it if it is
|
||||
# required.
|
||||
|
||||
batch_end = []
|
||||
responses = []
|
||||
batch_corrid = []
|
||||
qualified_corrid = []
|
||||
batch_result = {}
|
||||
inference_response_awaits = []
|
||||
|
||||
for request in requests:
|
||||
hidden = pb_utils.get_input_tensor_by_name(request, "enc")
|
||||
hidden = from_dlpack(hidden.to_dlpack()).cpu().numpy()
|
||||
alphas = pb_utils.get_input_tensor_by_name(request, "alphas")
|
||||
alphas = from_dlpack(alphas.to_dlpack()).cpu().numpy()
|
||||
hidden_len = pb_utils.get_input_tensor_by_name(request, "enc_len")
|
||||
hidden_len = from_dlpack(hidden_len.to_dlpack()).cpu().numpy()
|
||||
|
||||
in_start = pb_utils.get_input_tensor_by_name(request, "START")
|
||||
start = in_start.as_numpy()[0][0]
|
||||
|
||||
in_corrid = pb_utils.get_input_tensor_by_name(request, "CORRID")
|
||||
corrid = in_corrid.as_numpy()[0][0]
|
||||
|
||||
in_end = pb_utils.get_input_tensor_by_name(request, "END")
|
||||
end = in_end.as_numpy()[0][0]
|
||||
|
||||
batch_end.append(end)
|
||||
batch_corrid.append(corrid)
|
||||
|
||||
if start:
|
||||
self.cif_search_cache[corrid] = CIFSearch()
|
||||
self.start[corrid] = 1
|
||||
if end:
|
||||
self.cif_search_cache[corrid].cache["last_chunk"] = True
|
||||
|
||||
acoustic, acoustic_len = self.cif_search_cache[corrid].infer(hidden, alphas)
|
||||
batch_result[corrid] = ""
|
||||
if acoustic.shape[1] == 0:
|
||||
continue
|
||||
else:
|
||||
qualified_corrid.append(corrid)
|
||||
input_tensor0 = pb_utils.Tensor("enc", hidden)
|
||||
input_tensor1 = pb_utils.Tensor("enc_len", np.array([hidden_len], dtype=np.int32))
|
||||
input_tensor2 = pb_utils.Tensor("acoustic_embeds", acoustic)
|
||||
input_tensor3 = pb_utils.Tensor(
|
||||
"acoustic_embeds_len", np.array([acoustic_len], dtype=np.int32)
|
||||
)
|
||||
input_tensors = [input_tensor0, input_tensor1, input_tensor2, input_tensor3]
|
||||
|
||||
if self.start[corrid] and end:
|
||||
flag = 3
|
||||
elif end:
|
||||
flag = 2
|
||||
elif self.start[corrid]:
|
||||
flag = 1
|
||||
self.start[corrid] = 0
|
||||
else:
|
||||
flag = 0
|
||||
inference_request = pb_utils.InferenceRequest(
|
||||
model_name="decoder",
|
||||
requested_output_names=["sample_ids"],
|
||||
inputs=input_tensors,
|
||||
request_id="",
|
||||
correlation_id=corrid,
|
||||
flags=flag,
|
||||
)
|
||||
inference_response_awaits.append(inference_request.async_exec())
|
||||
|
||||
inference_responses = await asyncio.gather(*inference_response_awaits)
|
||||
|
||||
for index_corrid, inference_response in zip(qualified_corrid, inference_responses):
|
||||
if inference_response.has_error():
|
||||
raise pb_utils.TritonModelException(inference_response.error().message())
|
||||
else:
|
||||
sample_ids = pb_utils.get_output_tensor_by_name(inference_response, "sample_ids")
|
||||
token_ids = from_dlpack(sample_ids.to_dlpack()).cpu().numpy()[0]
|
||||
|
||||
# Change integer-ids to tokens
|
||||
tokens = [self.vocab_dict[token_id] for token_id in token_ids]
|
||||
batch_result[index_corrid] = "".join(tokens)
|
||||
|
||||
for i, index_corrid in enumerate(batch_corrid):
|
||||
sent = np.array([batch_result[index_corrid]])
|
||||
out0 = pb_utils.Tensor("transcripts", sent.astype(self.out0_dtype))
|
||||
inference_response = pb_utils.InferenceResponse(output_tensors=[out0])
|
||||
responses.append(inference_response)
|
||||
|
||||
if batch_end[i]:
|
||||
del self.cif_search_cache[index_corrid]
|
||||
del self.start[index_corrid]
|
||||
|
||||
return responses
|
||||
|
||||
def finalize(self):
|
||||
"""`finalize` is called only once when the model is being unloaded.
|
||||
Implementing `finalize` function is optional. This function allows
|
||||
the model to perform any necessary clean ups before exit.
|
||||
"""
|
||||
print("Cleaning up...")
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
name: "cif_search"
|
||||
backend: "python"
|
||||
max_batch_size: 128
|
||||
|
||||
sequence_batching{
|
||||
max_sequence_idle_microseconds: 15000000
|
||||
oldest {
|
||||
max_candidate_sequences: 1024
|
||||
preferred_batch_size: [32, 64, 128]
|
||||
}
|
||||
control_input [
|
||||
{
|
||||
name: "START",
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_START
|
||||
fp32_false_true: [0, 1]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
name: "READY"
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_READY
|
||||
fp32_false_true: [0, 1]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
name: "CORRID",
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_CORRID
|
||||
data_type: TYPE_UINT64
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
name: "END",
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_END
|
||||
fp32_false_true: [0, 1]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
parameters [
|
||||
{
|
||||
key: "vocabulary",
|
||||
value: { string_value: "model_repo_paraformer_large_online/feature_extractor/config.yaml"}
|
||||
},
|
||||
{ key: "FORCE_CPU_ONLY_INPUT_TENSORS"
|
||||
value: {string_value:"no"}
|
||||
}
|
||||
]
|
||||
|
||||
input [
|
||||
{
|
||||
name: "enc"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 512]
|
||||
},
|
||||
{
|
||||
name: "enc_len"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
reshape: { shape: [ ] }
|
||||
},
|
||||
{
|
||||
name: 'alphas'
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1]
|
||||
}
|
||||
]
|
||||
|
||||
output [
|
||||
{
|
||||
name: "transcripts"
|
||||
data_type: TYPE_STRING
|
||||
dims: [1]
|
||||
}
|
||||
]
|
||||
|
||||
instance_group [
|
||||
{
|
||||
count: 6
|
||||
kind: KIND_CPU
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,274 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
name: "decoder"
|
||||
backend: "onnxruntime"
|
||||
default_model_filename: "decoder.onnx"
|
||||
|
||||
max_batch_size: 128
|
||||
|
||||
sequence_batching{
|
||||
max_sequence_idle_microseconds: 15000000
|
||||
oldest {
|
||||
max_candidate_sequences: 1024
|
||||
preferred_batch_size: [16, 32, 64]
|
||||
}
|
||||
control_input [
|
||||
]
|
||||
state [
|
||||
{
|
||||
input_name: "in_cache_0"
|
||||
output_name: "out_cache_0"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_1"
|
||||
output_name: "out_cache_1"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_2"
|
||||
output_name: "out_cache_2"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_3"
|
||||
output_name: "out_cache_3"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_4"
|
||||
output_name: "out_cache_4"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_5"
|
||||
output_name: "out_cache_5"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_6"
|
||||
output_name: "out_cache_6"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_7"
|
||||
output_name: "out_cache_7"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_8"
|
||||
output_name: "out_cache_8"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_9"
|
||||
output_name: "out_cache_9"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_10"
|
||||
output_name: "out_cache_10"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_11"
|
||||
output_name: "out_cache_11"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_12"
|
||||
output_name: "out_cache_12"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_13"
|
||||
output_name: "out_cache_13"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_14"
|
||||
output_name: "out_cache_14"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "in_cache_15"
|
||||
output_name: "out_cache_15"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10 ]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 512, 10]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
input [
|
||||
{
|
||||
name: "enc"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 512]
|
||||
},
|
||||
{
|
||||
name: "enc_len"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
reshape: { shape: [ ] }
|
||||
},
|
||||
{
|
||||
name: "acoustic_embeds"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 512]
|
||||
},
|
||||
{
|
||||
name: "acoustic_embeds_len"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
reshape: { shape: [ ] }
|
||||
}
|
||||
]
|
||||
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 8404]
|
||||
},
|
||||
{
|
||||
name: "sample_ids"
|
||||
data_type: TYPE_INT64
|
||||
dims: [-1]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
instance_group [
|
||||
{
|
||||
count: 1
|
||||
kind: KIND_GPU
|
||||
}
|
||||
]
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
name: "encoder"
|
||||
backend: "onnxruntime"
|
||||
default_model_filename: "model.onnx"
|
||||
|
||||
max_batch_size: 128
|
||||
|
||||
|
||||
sequence_batching{
|
||||
max_sequence_idle_microseconds: 15000000
|
||||
oldest {
|
||||
max_candidate_sequences: 1024
|
||||
preferred_batch_size: [32, 64, 128]
|
||||
max_queue_delay_microseconds: 300
|
||||
}
|
||||
control_input [
|
||||
]
|
||||
state [
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
input [
|
||||
{
|
||||
name: "speech"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 560]
|
||||
},
|
||||
{
|
||||
name: "speech_lengths"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
reshape: { shape: [ ] }
|
||||
}
|
||||
]
|
||||
|
||||
output [
|
||||
{
|
||||
name: "enc"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 512]
|
||||
},
|
||||
{
|
||||
name: "enc_len"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
reshape: { shape: [ ] }
|
||||
},
|
||||
{
|
||||
name: "alphas"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1]
|
||||
}
|
||||
]
|
||||
|
||||
instance_group [
|
||||
{
|
||||
count: 1
|
||||
kind: KIND_GPU
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,216 @@
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
# Modified from NVIDIA(https://github.com/wenet-e2e/wenet/blob/main/runtime/gpu/
|
||||
# model_repo_stateful/feature_extractor/1/model.py)
|
||||
|
||||
import triton_python_backend_utils as pb_utils
|
||||
from torch.utils.dlpack import from_dlpack
|
||||
import torch
|
||||
import kaldifeat
|
||||
from typing import List
|
||||
import json
|
||||
import numpy as np
|
||||
import yaml
|
||||
from collections import OrderedDict
|
||||
|
||||
|
||||
class LimitedDict(OrderedDict):
|
||||
def __init__(self, max_length):
|
||||
super().__init__()
|
||||
self.max_length = max_length
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if len(self) >= self.max_length:
|
||||
self.popitem(last=False)
|
||||
super().__setitem__(key, value)
|
||||
|
||||
|
||||
class Fbank(torch.nn.Module):
|
||||
def __init__(self, opts):
|
||||
super(Fbank, self).__init__()
|
||||
self.fbank = kaldifeat.Fbank(opts)
|
||||
|
||||
def forward(self, waves: List[torch.Tensor]):
|
||||
return self.fbank(waves)
|
||||
|
||||
|
||||
class Feat(object):
|
||||
def __init__(self, seqid, offset_ms, sample_rate, frame_stride, device="cpu"):
|
||||
self.seqid = seqid
|
||||
self.sample_rate = sample_rate
|
||||
self.wav = torch.tensor([], device=device)
|
||||
self.offset = int(offset_ms / 1000 * sample_rate)
|
||||
self.frames = None
|
||||
self.frame_stride = int(frame_stride)
|
||||
self.device = device
|
||||
self.lfr_m = 7
|
||||
|
||||
def add_wavs(self, wav: torch.tensor):
|
||||
wav = wav.to(self.device)
|
||||
self.wav = torch.cat((self.wav, wav), axis=0)
|
||||
|
||||
def get_seg_wav(self):
|
||||
seg = self.wav[:]
|
||||
self.wav = self.wav[-self.offset :]
|
||||
return seg
|
||||
|
||||
def add_frames(self, frames: torch.tensor):
|
||||
"""
|
||||
frames: seq_len x feat_sz
|
||||
"""
|
||||
if self.frames is None:
|
||||
self.frames = torch.cat((frames[0, :].repeat((self.lfr_m - 1) // 2, 1), frames), axis=0)
|
||||
else:
|
||||
self.frames = torch.cat([self.frames, frames], axis=0)
|
||||
|
||||
def get_frames(self, num_frames: int):
|
||||
seg = self.frames[0:num_frames]
|
||||
self.frames = self.frames[self.frame_stride :]
|
||||
return seg
|
||||
|
||||
|
||||
class TritonPythonModel:
|
||||
"""Your Python model must use the same class name. Every Python model
|
||||
that is created must have "TritonPythonModel" as the class name.
|
||||
"""
|
||||
|
||||
def initialize(self, args):
|
||||
"""`initialize` is called only once when the model is being loaded.
|
||||
Implementing `initialize` function is optional. This function allows
|
||||
the model to initialize any state associated with this model.
|
||||
Parameters
|
||||
----------
|
||||
args : dict
|
||||
Both keys and values are strings. The dictionary keys and values are:
|
||||
* model_config: A JSON string containing the model configuration
|
||||
* model_instance_kind: A string containing model instance kind
|
||||
* model_instance_device_id: A string containing model instance device ID
|
||||
* model_repository: Model repository path
|
||||
* model_version: Model version
|
||||
* model_name: Model name
|
||||
"""
|
||||
self.model_config = model_config = json.loads(args["model_config"])
|
||||
self.max_batch_size = max(model_config["max_batch_size"], 1)
|
||||
|
||||
if "GPU" in model_config["instance_group"][0]["kind"]:
|
||||
self.device = "cuda"
|
||||
else:
|
||||
self.device = "cpu"
|
||||
|
||||
# Get OUTPUT0 configuration
|
||||
output0_config = pb_utils.get_output_config_by_name(model_config, "speech")
|
||||
# Convert Triton types to numpy types
|
||||
self.output0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"])
|
||||
|
||||
if self.output0_dtype == np.float32:
|
||||
self.dtype = torch.float32
|
||||
else:
|
||||
self.dtype = torch.float16
|
||||
|
||||
self.feature_size = output0_config["dims"][-1]
|
||||
self.decoding_window = output0_config["dims"][-2]
|
||||
|
||||
params = self.model_config["parameters"]
|
||||
for li in params.items():
|
||||
key, value = li
|
||||
value = value["string_value"]
|
||||
if key == "config_path":
|
||||
with open(str(value), "rb") as f:
|
||||
config = yaml.load(f, Loader=yaml.Loader)
|
||||
|
||||
opts = kaldifeat.FbankOptions()
|
||||
opts.frame_opts.dither = 0.0
|
||||
opts.frame_opts.window_type = config["frontend_conf"]["window"]
|
||||
opts.mel_opts.num_bins = int(config["frontend_conf"]["n_mels"])
|
||||
opts.frame_opts.frame_shift_ms = float(config["frontend_conf"]["frame_shift"])
|
||||
opts.frame_opts.frame_length_ms = float(config["frontend_conf"]["frame_length"])
|
||||
opts.frame_opts.samp_freq = int(config["frontend_conf"]["fs"])
|
||||
opts.device = torch.device(self.device)
|
||||
self.opts = opts
|
||||
self.feature_extractor = Fbank(self.opts)
|
||||
|
||||
self.seq_feat = LimitedDict(1024)
|
||||
chunk_size_s = float(params["chunk_size_s"]["string_value"])
|
||||
|
||||
sample_rate = opts.frame_opts.samp_freq
|
||||
frame_shift_ms = opts.frame_opts.frame_shift_ms
|
||||
frame_length_ms = opts.frame_opts.frame_length_ms
|
||||
|
||||
self.chunk_size = int(chunk_size_s * sample_rate)
|
||||
self.frame_stride = (chunk_size_s * 1000) // frame_shift_ms
|
||||
self.offset_ms = self.get_offset(frame_length_ms, frame_shift_ms)
|
||||
self.sample_rate = sample_rate
|
||||
|
||||
def get_offset(self, frame_length_ms, frame_shift_ms):
|
||||
offset_ms = 0
|
||||
while offset_ms + frame_shift_ms < frame_length_ms:
|
||||
offset_ms += frame_shift_ms
|
||||
return offset_ms
|
||||
|
||||
def execute(self, requests):
|
||||
"""`execute` must be implemented in every Python model. `execute`
|
||||
function receives a list of pb_utils.InferenceRequest as the only
|
||||
argument. This function is called when an inference is requested
|
||||
for this model.
|
||||
Parameters
|
||||
----------
|
||||
requests : list
|
||||
A list of pb_utils.InferenceRequest
|
||||
Returns
|
||||
-------
|
||||
list
|
||||
A list of pb_utils.InferenceResponse. The length of this list must
|
||||
be the same as `requests`
|
||||
"""
|
||||
total_waves = []
|
||||
responses = []
|
||||
batch_seqid = []
|
||||
end_seqid = {}
|
||||
for request in requests:
|
||||
input0 = pb_utils.get_input_tensor_by_name(request, "wav")
|
||||
wav = from_dlpack(input0.to_dlpack())[0]
|
||||
# input1 = pb_utils.get_input_tensor_by_name(request, "wav_lens")
|
||||
# wav_len = from_dlpack(input1.to_dlpack())[0]
|
||||
wav_len = len(wav)
|
||||
if wav_len < self.chunk_size:
|
||||
temp = torch.zeros(self.chunk_size, dtype=torch.float32, device=self.device)
|
||||
temp[0:wav_len] = wav[:]
|
||||
wav = temp
|
||||
|
||||
in_start = pb_utils.get_input_tensor_by_name(request, "START")
|
||||
start = in_start.as_numpy()[0][0]
|
||||
in_ready = pb_utils.get_input_tensor_by_name(request, "READY")
|
||||
ready = in_ready.as_numpy()[0][0]
|
||||
in_corrid = pb_utils.get_input_tensor_by_name(request, "CORRID")
|
||||
corrid = in_corrid.as_numpy()[0][0]
|
||||
in_end = pb_utils.get_input_tensor_by_name(request, "END")
|
||||
end = in_end.as_numpy()[0][0]
|
||||
|
||||
if start:
|
||||
self.seq_feat[corrid] = Feat(
|
||||
corrid, self.offset_ms, self.sample_rate, self.frame_stride, self.device
|
||||
)
|
||||
if ready:
|
||||
self.seq_feat[corrid].add_wavs(wav)
|
||||
|
||||
batch_seqid.append(corrid)
|
||||
if end:
|
||||
end_seqid[corrid] = 1
|
||||
|
||||
wav = self.seq_feat[corrid].get_seg_wav() * 32768
|
||||
total_waves.append(wav)
|
||||
features = self.feature_extractor(total_waves)
|
||||
for corrid, frames in zip(batch_seqid, features):
|
||||
self.seq_feat[corrid].add_frames(frames)
|
||||
speech = self.seq_feat[corrid].get_frames(self.decoding_window)
|
||||
out_tensor0 = pb_utils.Tensor("speech", torch.unsqueeze(speech, 0).to("cpu").numpy())
|
||||
output_tensors = [out_tensor0]
|
||||
response = pb_utils.InferenceResponse(output_tensors=output_tensors)
|
||||
responses.append(response)
|
||||
if corrid in end_seqid:
|
||||
del self.seq_feat[corrid]
|
||||
return responses
|
||||
|
||||
def finalize(self):
|
||||
print("Remove feature extractor!")
|
||||
@@ -0,0 +1,109 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
name: "feature_extractor"
|
||||
backend: "python"
|
||||
max_batch_size: 128
|
||||
|
||||
parameters [
|
||||
{
|
||||
key: "chunk_size_s",
|
||||
value: { string_value: "0.6"}
|
||||
},
|
||||
{
|
||||
key: "config_path"
|
||||
value: { string_value: "model_repo_paraformer_large_online/feature_extractor/config.yaml"}
|
||||
}
|
||||
]
|
||||
|
||||
sequence_batching{
|
||||
max_sequence_idle_microseconds: 15000000
|
||||
oldest {
|
||||
max_candidate_sequences: 1024
|
||||
preferred_batch_size: [32, 64, 128]
|
||||
max_queue_delay_microseconds: 300
|
||||
}
|
||||
control_input [
|
||||
{
|
||||
name: "START",
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_START
|
||||
fp32_false_true: [0, 1]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
name: "READY"
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_READY
|
||||
fp32_false_true: [0, 1]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
name: "CORRID",
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_CORRID
|
||||
data_type: TYPE_UINT64
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
name: "END",
|
||||
control [
|
||||
{
|
||||
kind: CONTROL_SEQUENCE_END
|
||||
fp32_false_true: [0, 1]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
input [
|
||||
{
|
||||
name: "wav"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1]
|
||||
},
|
||||
{
|
||||
name: "wav_lens"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
}
|
||||
]
|
||||
|
||||
output [
|
||||
{
|
||||
name: "speech"
|
||||
data_type: TYPE_FP32
|
||||
dims: [61, 80] # 80
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
instance_group [
|
||||
{
|
||||
count: 1
|
||||
kind: KIND_GPU
|
||||
}
|
||||
]
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,8 @@
|
||||
<Nnet>
|
||||
<Splice> 560 560
|
||||
[ 0 ]
|
||||
<AddShift> 560 560
|
||||
<LearnRateCoef> 0 [ -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.65519 -13.50713 -13.75251 -13.76715 -13.87408 -13.73109 -13.70412 -13.56073 -13.53488 -13.54895 -13.56228 -13.59408 -13.62047 -13.64198 -13.66109 -13.62669 -13.58297 -13.57387 -13.4739 -13.53063 -13.48348 -13.61047 -13.64716 -13.71546 -13.79184 -13.90614 -14.03098 -14.18205 -14.35881 -14.48419 -14.60172 -14.70591 -14.83362 -14.92122 -15.00622 -15.05122 -15.03119 -14.99028 -14.92302 -14.86927 -14.82691 -14.7972 -14.76909 -14.71356 -14.61277 -14.51696 -14.42252 -14.36405 -14.30451 -14.23161 -14.19851 -14.16633 -14.15649 -14.10504 -13.99518 -13.79562 -13.3996 -12.7767 -11.71208 -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.65519 -13.50713 -13.75251 -13.76715 -13.87408 -13.73109 -13.70412 -13.56073 -13.53488 -13.54895 -13.56228 -13.59408 -13.62047 -13.64198 -13.66109 -13.62669 -13.58297 -13.57387 -13.4739 -13.53063 -13.48348 -13.61047 -13.64716 -13.71546 -13.79184 -13.90614 -14.03098 -14.18205 -14.35881 -14.48419 -14.60172 -14.70591 -14.83362 -14.92122 -15.00622 -15.05122 -15.03119 -14.99028 -14.92302 -14.86927 -14.82691 -14.7972 -14.76909 -14.71356 -14.61277 -14.51696 -14.42252 -14.36405 -14.30451 -14.23161 -14.19851 -14.16633 -14.15649 -14.10504 -13.99518 -13.79562 -13.3996 -12.7767 -11.71208 -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.65519 -13.50713 -13.75251 -13.76715 -13.87408 -13.73109 -13.70412 -13.56073 -13.53488 -13.54895 -13.56228 -13.59408 -13.62047 -13.64198 -13.66109 -13.62669 -13.58297 -13.57387 -13.4739 -13.53063 -13.48348 -13.61047 -13.64716 -13.71546 -13.79184 -13.90614 -14.03098 -14.18205 -14.35881 -14.48419 -14.60172 -14.70591 -14.83362 -14.92122 -15.00622 -15.05122 -15.03119 -14.99028 -14.92302 -14.86927 -14.82691 -14.7972 -14.76909 -14.71356 -14.61277 -14.51696 -14.42252 -14.36405 -14.30451 -14.23161 -14.19851 -14.16633 -14.15649 -14.10504 -13.99518 -13.79562 -13.3996 -12.7767 -11.71208 -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.65519 -13.50713 -13.75251 -13.76715 -13.87408 -13.73109 -13.70412 -13.56073 -13.53488 -13.54895 -13.56228 -13.59408 -13.62047 -13.64198 -13.66109 -13.62669 -13.58297 -13.57387 -13.4739 -13.53063 -13.48348 -13.61047 -13.64716 -13.71546 -13.79184 -13.90614 -14.03098 -14.18205 -14.35881 -14.48419 -14.60172 -14.70591 -14.83362 -14.92122 -15.00622 -15.05122 -15.03119 -14.99028 -14.92302 -14.86927 -14.82691 -14.7972 -14.76909 -14.71356 -14.61277 -14.51696 -14.42252 -14.36405 -14.30451 -14.23161 -14.19851 -14.16633 -14.15649 -14.10504 -13.99518 -13.79562 -13.3996 -12.7767 -11.71208 -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.65519 -13.50713 -13.75251 -13.76715 -13.87408 -13.73109 -13.70412 -13.56073 -13.53488 -13.54895 -13.56228 -13.59408 -13.62047 -13.64198 -13.66109 -13.62669 -13.58297 -13.57387 -13.4739 -13.53063 -13.48348 -13.61047 -13.64716 -13.71546 -13.79184 -13.90614 -14.03098 -14.18205 -14.35881 -14.48419 -14.60172 -14.70591 -14.83362 -14.92122 -15.00622 -15.05122 -15.03119 -14.99028 -14.92302 -14.86927 -14.82691 -14.7972 -14.76909 -14.71356 -14.61277 -14.51696 -14.42252 -14.36405 -14.30451 -14.23161 -14.19851 -14.16633 -14.15649 -14.10504 -13.99518 -13.79562 -13.3996 -12.7767 -11.71208 -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.65519 -13.50713 -13.75251 -13.76715 -13.87408 -13.73109 -13.70412 -13.56073 -13.53488 -13.54895 -13.56228 -13.59408 -13.62047 -13.64198 -13.66109 -13.62669 -13.58297 -13.57387 -13.4739 -13.53063 -13.48348 -13.61047 -13.64716 -13.71546 -13.79184 -13.90614 -14.03098 -14.18205 -14.35881 -14.48419 -14.60172 -14.70591 -14.83362 -14.92122 -15.00622 -15.05122 -15.03119 -14.99028 -14.92302 -14.86927 -14.82691 -14.7972 -14.76909 -14.71356 -14.61277 -14.51696 -14.42252 -14.36405 -14.30451 -14.23161 -14.19851 -14.16633 -14.15649 -14.10504 -13.99518 -13.79562 -13.3996 -12.7767 -11.71208 -8.311879 -8.600912 -9.615928 -10.43595 -11.21292 -11.88333 -12.36243 -12.63706 -12.8818 -12.83066 -12.89103 -12.95666 -13.19763 -13.40598 -13.49113 -13.5546 -13.55639 -13.51915 -13.68284 -13.53289 -13.42107 -13.655Line truncated
|
||||
<Rescale> 560 560
|
||||
<LearnRateCoef> 0 [ 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 0.147664 0.1480913 0.1483739 0.1488841 0.1493636 0.1497088 0.1500379 0.1502916 0.1505389 0.1506787 0.1507102 0.1505992 0.1505445 0.1505938 0.1508133 0.1509569 0.1512396 0.1514625 0.1516195 0.1516156 0.1515561 0.1514966 0.1513976 0.1512612 0.151076 0.1510596 0.1510431 0.151077 0.1511168 0.1511917 0.151023 0.1508045 0.1505885 0.1503493 0.1502373 0.1501726 0.1500762 0.1500065 0.1499782 0.150057 0.1502658 0.150469 0.1505335 0.1505505 0.1505328 0.1504275 0.1502438 0.1499674 0.1497118 0.1494661 0.1493102 0.1493681 0.1495501 0.1499738 0.1509654 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 0.147664 0.1480913 0.1483739 0.1488841 0.1493636 0.1497088 0.1500379 0.1502916 0.1505389 0.1506787 0.1507102 0.1505992 0.1505445 0.1505938 0.1508133 0.1509569 0.1512396 0.1514625 0.1516195 0.1516156 0.1515561 0.1514966 0.1513976 0.1512612 0.151076 0.1510596 0.1510431 0.151077 0.1511168 0.1511917 0.151023 0.1508045 0.1505885 0.1503493 0.1502373 0.1501726 0.1500762 0.1500065 0.1499782 0.150057 0.1502658 0.150469 0.1505335 0.1505505 0.1505328 0.1504275 0.1502438 0.1499674 0.1497118 0.1494661 0.1493102 0.1493681 0.1495501 0.1499738 0.1509654 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 0.147664 0.1480913 0.1483739 0.1488841 0.1493636 0.1497088 0.1500379 0.1502916 0.1505389 0.1506787 0.1507102 0.1505992 0.1505445 0.1505938 0.1508133 0.1509569 0.1512396 0.1514625 0.1516195 0.1516156 0.1515561 0.1514966 0.1513976 0.1512612 0.151076 0.1510596 0.1510431 0.151077 0.1511168 0.1511917 0.151023 0.1508045 0.1505885 0.1503493 0.1502373 0.1501726 0.1500762 0.1500065 0.1499782 0.150057 0.1502658 0.150469 0.1505335 0.1505505 0.1505328 0.1504275 0.1502438 0.1499674 0.1497118 0.1494661 0.1493102 0.1493681 0.1495501 0.1499738 0.1509654 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 0.147664 0.1480913 0.1483739 0.1488841 0.1493636 0.1497088 0.1500379 0.1502916 0.1505389 0.1506787 0.1507102 0.1505992 0.1505445 0.1505938 0.1508133 0.1509569 0.1512396 0.1514625 0.1516195 0.1516156 0.1515561 0.1514966 0.1513976 0.1512612 0.151076 0.1510596 0.1510431 0.151077 0.1511168 0.1511917 0.151023 0.1508045 0.1505885 0.1503493 0.1502373 0.1501726 0.1500762 0.1500065 0.1499782 0.150057 0.1502658 0.150469 0.1505335 0.1505505 0.1505328 0.1504275 0.1502438 0.1499674 0.1497118 0.1494661 0.1493102 0.1493681 0.1495501 0.1499738 0.1509654 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 0.147664 0.1480913 0.1483739 0.1488841 0.1493636 0.1497088 0.1500379 0.1502916 0.1505389 0.1506787 0.1507102 0.1505992 0.1505445 0.1505938 0.1508133 0.1509569 0.1512396 0.1514625 0.1516195 0.1516156 0.1515561 0.1514966 0.1513976 0.1512612 0.151076 0.1510596 0.1510431 0.151077 0.1511168 0.1511917 0.151023 0.1508045 0.1505885 0.1503493 0.1502373 0.1501726 0.1500762 0.1500065 0.1499782 0.150057 0.1502658 0.150469 0.1505335 0.1505505 0.1505328 0.1504275 0.1502438 0.1499674 0.1497118 0.1494661 0.1493102 0.1493681 0.1495501 0.1499738 0.1509654 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 0.147664 0.1480913 0.1483739 0.1488841 0.1493636 0.1497088 0.1500379 0.1502916 0.1505389 0.1506787 0.1507102 0.1505992 0.1505445 0.1505938 0.1508133 0.1509569 0.1512396 0.1514625 0.1516195 0.1516156 0.1515561 0.1514966 0.1513976 0.1512612 0.151076 0.1510596 0.1510431 0.151077 0.1511168 0.1511917 0.151023 0.1508045 0.1505885 0.1503493 0.1502373 0.1501726 0.1500762 0.1500065 0.1499782 0.150057 0.1502658 0.150469 0.1505335 0.1505505 0.1505328 0.1504275 0.1502438 0.1499674 0.1497118 0.1494661 0.1493102 0.1493681 0.1495501 0.1499738 0.1509654 0.155775 0.154484 0.1527379 0.1518718 0.1506028 0.1489256 0.147067 0.1447061 0.1436307 0.1443568 0.1451849 0.1455157 0.1452821 0.1445717 0.1439195 0.1435867 0.1436018 0.1438781 0.1442086 0.1448844 0.1454756 0.145663 0.146268 0.1467386 0.1472724 Line truncated
|
||||
</Nnet>
|
||||
@@ -0,0 +1,85 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
name: "lfr_cmvn_pe"
|
||||
backend: "onnxruntime"
|
||||
default_model_filename: "lfr_cmvn_pe.onnx"
|
||||
|
||||
max_batch_size: 128
|
||||
|
||||
sequence_batching{
|
||||
max_sequence_idle_microseconds: 15000000
|
||||
oldest {
|
||||
max_candidate_sequences: 1024
|
||||
preferred_batch_size: [32, 64, 128]
|
||||
max_queue_delay_microseconds: 300
|
||||
}
|
||||
control_input [
|
||||
]
|
||||
state [
|
||||
{
|
||||
input_name: "cache"
|
||||
output_name: "r_cache"
|
||||
data_type: TYPE_FP32
|
||||
dims: [10, 560]
|
||||
initial_state: {
|
||||
data_type: TYPE_FP32
|
||||
dims: [10, 560]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
},
|
||||
{
|
||||
input_name: "offset"
|
||||
output_name: "r_offset"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
initial_state: {
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
zero_data: true
|
||||
name: "initial state"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
input [
|
||||
{
|
||||
name: "chunk_xs"
|
||||
data_type: TYPE_FP32
|
||||
dims: [61, 80]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "chunk_xs_out"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1, 560]
|
||||
},
|
||||
{
|
||||
name: "chunk_xs_out_len"
|
||||
data_type: TYPE_INT32
|
||||
dims: [-1]
|
||||
}
|
||||
]
|
||||
instance_group [
|
||||
{
|
||||
count: 1
|
||||
kind: KIND_GPU
|
||||
}
|
||||
]
|
||||
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import math
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class LFR_CMVN_PE(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
mean: torch.Tensor,
|
||||
istd: torch.Tensor,
|
||||
m: int = 7,
|
||||
n: int = 6,
|
||||
max_len: int = 5000,
|
||||
encoder_input_size: int = 560,
|
||||
encoder_output_size: int = 512,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# LRF
|
||||
self.m = m
|
||||
self.n = n
|
||||
self.subsample = (m - 1) // 2
|
||||
|
||||
# CMVN
|
||||
assert mean.shape == istd.shape
|
||||
# The buffer can be accessed from this module using self.mean
|
||||
self.register_buffer("mean", mean)
|
||||
self.register_buffer("istd", istd)
|
||||
|
||||
# PE
|
||||
self.encoder_input_size = encoder_input_size
|
||||
self.encoder_output_size = encoder_output_size
|
||||
self.max_len = max_len
|
||||
self.pe = torch.zeros(self.max_len, self.encoder_input_size)
|
||||
position = torch.arange(0, self.max_len, dtype=torch.float32).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange((self.encoder_input_size / 2), dtype=torch.float32)
|
||||
* -(math.log(10000.0) / (self.encoder_input_size / 2 - 1))
|
||||
)
|
||||
self.pe[:, 0::1] = torch.cat(
|
||||
(torch.sin(position * div_term), torch.cos(position * div_term)), dim=1
|
||||
)
|
||||
|
||||
def forward(self, x, cache, offset):
|
||||
"""
|
||||
Args:
|
||||
x (torch.Tensor): (batch, max_len, feat_dim)
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): normalized feature
|
||||
"""
|
||||
B, _, D = x.size()
|
||||
x = x.unfold(1, self.m, step=self.n).transpose(2, 3)
|
||||
x = x.view(B, -1, D * self.m)
|
||||
|
||||
x = (x + self.mean) * self.istd
|
||||
x = x * (self.encoder_output_size**0.5)
|
||||
|
||||
index = offset + torch.arange(1, x.size(1) + 1).to(dtype=torch.int32)
|
||||
pos_emb = F.embedding(index, self.pe) # B X T X d_model
|
||||
r_cache = x + pos_emb
|
||||
|
||||
r_x = torch.cat((cache, r_cache), dim=1)
|
||||
r_offset = offset + x.size(1)
|
||||
r_x_len = torch.ones((B, 1), dtype=torch.int32) * r_x.size(1)
|
||||
|
||||
return r_x, r_x_len, r_cache, r_offset
|
||||
|
||||
|
||||
def load_cmvn(cmvn_file):
|
||||
with open(cmvn_file, "r", encoding="utf-8") as f:
|
||||
lines = f.readlines()
|
||||
|
||||
means_list = []
|
||||
vars_list = []
|
||||
for i in range(len(lines)):
|
||||
line_item = lines[i].split()
|
||||
if line_item[0] == "<AddShift>":
|
||||
line_item = lines[i + 1].split()
|
||||
if line_item[0] == "<LearnRateCoef>":
|
||||
add_shift_line = line_item[3 : (len(line_item) - 1)]
|
||||
means_list = list(add_shift_line)
|
||||
continue
|
||||
elif line_item[0] == "<Rescale>":
|
||||
line_item = lines[i + 1].split()
|
||||
if line_item[0] == "<LearnRateCoef>":
|
||||
rescale_line = line_item[3 : (len(line_item) - 1)]
|
||||
vars_list = list(rescale_line)
|
||||
continue
|
||||
|
||||
means = np.array(means_list).astype(np.float32)
|
||||
vars = np.array(vars_list).astype(np.float32)
|
||||
means = torch.from_numpy(means)
|
||||
vars = torch.from_numpy(vars)
|
||||
return means, vars
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
means, vars = load_cmvn("am.mvn")
|
||||
means = torch.tile(means, (10, 1))
|
||||
vars = torch.tile(vars, (10, 1))
|
||||
|
||||
model = LFR_CMVN_PE(means, vars)
|
||||
model.eval()
|
||||
|
||||
all_names = [
|
||||
"chunk_xs",
|
||||
"cache",
|
||||
"offset",
|
||||
"chunk_xs_out",
|
||||
"chunk_xs_out_len",
|
||||
"r_cache",
|
||||
"r_offset",
|
||||
]
|
||||
dynamic_axes = {}
|
||||
|
||||
for name in all_names:
|
||||
dynamic_axes[name] = {0: "B"}
|
||||
|
||||
input_data1 = torch.randn(4, 61, 80).to(torch.float32)
|
||||
input_data2 = torch.randn(4, 10, 560).to(torch.float32)
|
||||
input_data3 = torch.randn(4, 1).to(torch.int32)
|
||||
|
||||
onnx_path = "./1/lfr_cmvn_pe.onnx"
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(input_data1, input_data2, input_data3),
|
||||
onnx_path,
|
||||
export_params=True,
|
||||
opset_version=11,
|
||||
do_constant_folding=True,
|
||||
input_names=["chunk_xs", "cache", "offset"],
|
||||
output_names=["chunk_xs_out", "chunk_xs_out_len", "r_cache", "r_offset"],
|
||||
dynamic_axes=dynamic_axes,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
print("export to onnx model succeed!")
|
||||
+122
@@ -0,0 +1,122 @@
|
||||
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Created on 2024-01-01
|
||||
# Author: GuAn Zhu
|
||||
|
||||
name: "streaming_paraformer"
|
||||
platform: "ensemble"
|
||||
max_batch_size: 128 #MAX_BATCH
|
||||
|
||||
input [
|
||||
{
|
||||
name: "WAV"
|
||||
data_type: TYPE_FP32
|
||||
dims: [-1]
|
||||
},
|
||||
{
|
||||
name: "WAV_LENS"
|
||||
data_type: TYPE_INT32
|
||||
dims: [1]
|
||||
}
|
||||
]
|
||||
|
||||
output [
|
||||
{
|
||||
name: "TRANSCRIPTS"
|
||||
data_type: TYPE_STRING
|
||||
dims: [1]
|
||||
}
|
||||
]
|
||||
|
||||
ensemble_scheduling {
|
||||
step [
|
||||
{
|
||||
model_name: "feature_extractor"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "wav"
|
||||
value: "WAV"
|
||||
}
|
||||
input_map {
|
||||
key: "wav_lens"
|
||||
value: "WAV_LENS"
|
||||
}
|
||||
output_map {
|
||||
key: "speech"
|
||||
value: "SPEECH"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "lfr_cmvn_pe"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "chunk_xs"
|
||||
value: "SPEECH"
|
||||
}
|
||||
output_map {
|
||||
key: "chunk_xs_out"
|
||||
value: "CHUNK_XS_OUT"
|
||||
}
|
||||
output_map {
|
||||
key: "chunk_xs_out_len"
|
||||
value: "CHUNK_XS_OUT_LEN"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "encoder"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "speech"
|
||||
value: "CHUNK_XS_OUT"
|
||||
}
|
||||
input_map {
|
||||
key: "speech_lengths"
|
||||
value: "CHUNK_XS_OUT_LEN"
|
||||
}
|
||||
output_map {
|
||||
key: "enc"
|
||||
value: "ENC"
|
||||
}
|
||||
output_map {
|
||||
key: "enc_len"
|
||||
value: "ENC_LEN"
|
||||
}
|
||||
output_map {
|
||||
key: "alphas"
|
||||
value: "ALPHAS"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "cif_search"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "enc"
|
||||
value: "ENC"
|
||||
}
|
||||
input_map {
|
||||
key: "enc_len"
|
||||
value: "ENC_LEN"
|
||||
}
|
||||
input_map {
|
||||
key: "alphas"
|
||||
value: "ALPHAS"
|
||||
}
|
||||
output_map {
|
||||
key: "transcripts"
|
||||
value: "TRANSCRIPTS"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
Reference in new issue
Block a user