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

Support batch greedy search decoding (#30)

Browse Source
This commit is contained in:
Fangjun Kuang
2023-02-19 15:04:24 +08:00
committed by GitHub
parent ebc3b47fb8
commit 8acc059b3f
5 changed files with 181 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
sherpa-onnx/csrc/online-transducer-greedy-search-decoder.cc
View File

@@ -32,6 +32,30 @@ static Ort::Value GetFrame(Ort::Value *encoder_out, int32_t t) {
encoder_out_dim, shape.data(), shape.size());
}
static Ort::Value Repeat(OrtAllocator *allocator, Ort::Value *cur_encoder_out,
int32_t n) {
if (n == 1) {
return std::move(*cur_encoder_out);
}
std::vector<int64_t> cur_encoder_out_shape =
cur_encoder_out->GetTensorTypeAndShapeInfo().GetShape();
std::array<int64_t, 2> ans_shape{n, cur_encoder_out_shape[1]};
Ort::Value ans = Ort::Value::CreateTensor<float>(allocator, ans_shape.data(),
ans_shape.size());
const float *src = cur_encoder_out->GetTensorData<float>();
float *dst = ans.GetTensorMutableData<float>();
for (int32_t i = 0; i != n; ++i) {
std::copy(src, src + cur_encoder_out_shape[1], dst);
dst += cur_encoder_out_shape[1];
}
return ans;
}
OnlineTransducerDecoderResult
OnlineTransducerGreedySearchDecoder::GetEmptyResult() const {
int32_t context_size = model_->ContextSize();
@@ -66,33 +90,33 @@ void OnlineTransducerGreedySearchDecoder::Decode(
exit(-1);
}
if (result->size() != 1) {
fprintf(stderr, "only batch size == 1 is implemented. Given: %d",
static_cast<int32_t>(result->size()));
exit(-1);
}
auto &hyp = (*result)[0].tokens;
int32_t num_frames = encoder_out_shape[1];
int32_t batch_size = static_cast<int32_t>(encoder_out_shape[0]);
int32_t num_frames = static_cast<int32_t>(encoder_out_shape[1]);
int32_t vocab_size = model_->VocabSize();
Ort::Value decoder_input = model_->BuildDecoderInput(hyp);
Ort::Value decoder_input = model_->BuildDecoderInput(*result);
Ort::Value decoder_out = model_->RunDecoder(std::move(decoder_input));
for (int32_t t = 0; t != num_frames; ++t) {
Ort::Value cur_encoder_out = GetFrame(&encoder_out, t);
cur_encoder_out = Repeat(model_->Allocator(), &cur_encoder_out, batch_size);
Ort::Value logit =
model_->RunJoiner(std::move(cur_encoder_out), Clone(&decoder_out));
const float *p_logit = logit.GetTensorData<float>();
auto y = static_cast<int32_t>(std::distance(
static_cast<const float *>(p_logit),
std::max_element(static_cast<const float *>(p_logit),
static_cast<const float *>(p_logit) + vocab_size)));
if (y != 0) {
hyp.push_back(y);
decoder_input = model_->BuildDecoderInput(hyp);
bool emitted = false;
for (int32_t i = 0; i < batch_size; ++i, p_logit += vocab_size) {
auto y = static_cast<int32_t>(std::distance(
static_cast<const float *>(p_logit),
std::max_element(static_cast<const float *>(p_logit),
static_cast<const float *>(p_logit) + vocab_size)));
if (y != 0) {
emitted = true;
(*result)[i].tokens.push_back(y);
}
}
if (emitted) {
decoder_input = model_->BuildDecoderInput(*result);
decoder_out = model_->RunDecoder(std::move(decoder_input));
}
}