EngineX-Iluvatar/enginex-mr_series-sherpa-onnx
Archived
16
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Add Streaming zipformer (#50)

Browse Source
This commit is contained in:
Fangjun Kuang
2023-02-21 20:00:03 +08:00
committed by GitHub
parent 8d1be0945e
commit 3ea6aa949d
20 changed files with 1576 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
sherpa-onnx/csrc/online-lstm-transducer-model.cc
View File

@@ -3,6 +3,8 @@
// Copyright (c) 2023 Xiaomi Corporation
#include "sherpa-onnx/csrc/online-lstm-transducer-model.h"
#include <assert.h>
#include <algorithm>
#include <memory>
#include <sstream>
@@ -11,23 +13,11 @@
#include <vector>
#include "onnxruntime_cxx_api.h" // NOLINT
#include "sherpa-onnx/csrc/cat.h"
#include "sherpa-onnx/csrc/macros.h"
#include "sherpa-onnx/csrc/online-transducer-decoder.h"
#include "sherpa-onnx/csrc/onnx-utils.h"
#define SHERPA_ONNX_READ_META_DATA(dst, src_key) \
do { \
auto value = \
meta_data.LookupCustomMetadataMapAllocated(src_key, allocator); \
if (!value) { \
fprintf(stderr, "%s does not exist in the metadata\n", src_key); \
exit(-1); \
} \
dst = atoi(value.get()); \
if (dst <= 0) { \
fprintf(stderr, "Invalud value %d for %s\n", dst, src_key); \
exit(-1); \
} \
} while (0)
#include "sherpa-onnx/csrc/unbind.h"
namespace sherpa_onnx {
@@ -64,7 +54,7 @@ void OnlineLstmTransducerModel::InitEncoder(const std::string &filename) {
fprintf(stderr, "%s\n", os.str().c_str());
}
Ort::AllocatorWithDefaultOptions allocator;
Ort::AllocatorWithDefaultOptions allocator; // used in the macro below
SHERPA_ONNX_READ_META_DATA(num_encoder_layers_, "num_encoder_layers");
SHERPA_ONNX_READ_META_DATA(T_, "T");
SHERPA_ONNX_READ_META_DATA(decode_chunk_len_, "decode_chunk_len");
@@ -91,7 +81,7 @@ void OnlineLstmTransducerModel::InitDecoder(const std::string &filename) {
fprintf(stderr, "%s\n", os.str().c_str());
}
Ort::AllocatorWithDefaultOptions allocator;
Ort::AllocatorWithDefaultOptions allocator; // used in the macro below
SHERPA_ONNX_READ_META_DATA(vocab_size_, "vocab_size");
SHERPA_ONNX_READ_META_DATA(context_size_, "context_size");
}
@@ -120,37 +110,19 @@ std::vector<Ort::Value> OnlineLstmTransducerModel::StackStates(
const std::vector<std::vector<Ort::Value>> &states) const {
int32_t batch_size = static_cast<int32_t>(states.size());
std::array<int64_t, 3> h_shape{num_encoder_layers_, batch_size, d_model_};
Ort::Value h = Ort::Value::CreateTensor<float>(allocator_, h_shape.data(),
h_shape.size());
std::vector<const Ort::Value *> h_buf(batch_size);
std::vector<const Ort::Value *> c_buf(batch_size);
std::array<int64_t, 3> c_shape{num_encoder_layers_, batch_size,
rnn_hidden_size_};
Ort::Value c = Ort::Value::CreateTensor<float>(allocator_, c_shape.data(),
c_shape.size());
float *dst_h = h.GetTensorMutableData<float>();
float *dst_c = c.GetTensorMutableData<float>();
for (int32_t layer = 0; layer != num_encoder_layers_; ++layer) {
for (int32_t i = 0; i != batch_size; ++i) {
const float *src_h =
states[i][0].GetTensorData<float>() + layer * d_model_;
const float *src_c =
states[i][1].GetTensorData<float>() + layer * rnn_hidden_size_;
std::copy(src_h, src_h + d_model_, dst_h);
std::copy(src_c, src_c + rnn_hidden_size_, dst_c);
dst_h += d_model_;
dst_c += rnn_hidden_size_;
}
for (int32_t i = 0; i != batch_size; ++i) {
assert(states[i].size() == 2);
h_buf[i] = &states[i][0];
c_buf[i] = &states[i][1];
}
std::vector<Ort::Value> ans;
Ort::Value h = Cat(allocator_, h_buf, 1);
Ort::Value c = Cat(allocator_, c_buf, 1);
std::vector<Ort::Value> ans;
ans.reserve(2);
ans.push_back(std::move(h));
ans.push_back(std::move(c));
@@ -161,37 +133,19 @@ std::vector<Ort::Value> OnlineLstmTransducerModel::StackStates(
std::vector<std::vector<Ort::Value>> OnlineLstmTransducerModel::UnStackStates(
const std::vector<Ort::Value> &states) const {
int32_t batch_size = states[0].GetTensorTypeAndShapeInfo().GetShape()[1];
assert(states.size() == 2);
std::vector<std::vector<Ort::Value>> ans(batch_size);
// allocate space
std::array<int64_t, 3> h_shape{num_encoder_layers_, 1, d_model_};
std::array<int64_t, 3> c_shape{num_encoder_layers_, 1, rnn_hidden_size_};
std::vector<Ort::Value> h_vec = Unbind(allocator_, &states[0], 1);
std::vector<Ort::Value> c_vec = Unbind(allocator_, &states[1], 1);
assert(h_vec.size() == batch_size);
assert(c_vec.size() == batch_size);
for (int32_t i = 0; i != batch_size; ++i) {
Ort::Value h = Ort::Value::CreateTensor<float>(allocator_, h_shape.data(),
h_shape.size());
Ort::Value c = Ort::Value::CreateTensor<float>(allocator_, c_shape.data(),
c_shape.size());
ans[i].push_back(std::move(h));
ans[i].push_back(std::move(c));
}
for (int32_t layer = 0; layer != num_encoder_layers_; ++layer) {
for (int32_t i = 0; i != batch_size; ++i) {
const float *src_h = states[0].GetTensorData<float>() +
layer * batch_size * d_model_ + i * d_model_;
const float *src_c = states[1].GetTensorData<float>() +
layer * batch_size * rnn_hidden_size_ +
i * rnn_hidden_size_;
float *dst_h = ans[i][0].GetTensorMutableData<float>() + layer * d_model_;
float *dst_c =
ans[i][1].GetTensorMutableData<float>() + layer * rnn_hidden_size_;
std::copy(src_h, src_h + d_model_, dst_h);
std::copy(src_c, src_c + rnn_hidden_size_, dst_c);
}
ans[i].push_back(std::move(h_vec[i]));
ans[i].push_back(std::move(c_vec[i]));
}
return ans;
@@ -206,20 +160,15 @@ std::vector<Ort::Value> OnlineLstmTransducerModel::GetEncoderInitStates() {
Ort::Value h = Ort::Value::CreateTensor<float>(allocator_, h_shape.data(),
h_shape.size());
std::fill(h.GetTensorMutableData<float>(),
h.GetTensorMutableData<float>() +
num_encoder_layers_ * kBatchSize * d_model_,
0);
Fill<float>(&h, 0);
std::array<int64_t, 3> c_shape{num_encoder_layers_, kBatchSize,
rnn_hidden_size_};
Ort::Value c = Ort::Value::CreateTensor<float>(allocator_, c_shape.data(),
c_shape.size());
std::fill(c.GetTensorMutableData<float>(),
c.GetTensorMutableData<float>() +
num_encoder_layers_ * kBatchSize * rnn_hidden_size_,
0);
Fill<float>(&c, 0);
std::vector<Ort::Value> states;