Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

runtime/triton_gpu/model_repo_paraformer_large_online/cif_search/1/model.py

@@ -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...")

runtime/triton_gpu/model_repo_paraformer_large_online/cif_search/config.pbtxt

@@ -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
+    }
+  ]

runtime/triton_gpu/model_repo_paraformer_large_online/decoder/config.pbtxt

@@ -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
+    }
+]
+

runtime/triton_gpu/model_repo_paraformer_large_online/encoder/config.pbtxt

@@ -0,0 +1,79 @@
+# 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]
+  }
+]
+
+parameters { key: "cudnn_conv_algo_search" value: { string_value: "2" } }
+
+instance_group [
+    {
+      count: 1
+      kind: KIND_GPU
+    }
+]

runtime/triton_gpu/model_repo_paraformer_large_online/feature_extractor/1/model.py

@@ -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!")

runtime/triton_gpu/model_repo_paraformer_large_online/feature_extractor/config.pbtxt

@@ -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
+    }
+]
+

runtime/triton_gpu/model_repo_paraformer_large_online/lfr_cmvn_pe/config.pbtxt

@@ -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
+    }
+]
+

runtime/triton_gpu/model_repo_paraformer_large_online/lfr_cmvn_pe/export_lfr_cmvn_pe_onnx.py

@@ -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!")

runtime/triton_gpu/model_repo_paraformer_large_online/streaming_paraformer/config.pbtxt

@@ -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"
+      }
+   }
+ ]
+}