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

Use log probs for paraformer (#120)

Browse Source 
* Use log probs for paraformer

* Fix
This commit is contained in:
彭震东
2023-04-09 23:28:10 +08:00
committed by GitHub
parent 80060c276d
commit d781fcdeef
3 changed files with 15 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
sherpa-onnx/csrc/offline-paraformer-decoder.h
View File

@@ -23,8 +23,7 @@ class OfflineParaformerDecoder {
/** Run beam search given the output from the paraformer model.
*
* @param log_probs A 3-D tensor of shape (N, T, vocab_size)
* @param token_num A 2-D tensor of shape (N, T). Its dtype is int64_t.
* log_probs[i].argmax(axis=-1) equals to token_num[i]
* @param token_num A 1-D tensor of shape (N). token_num equals to T.
*
* @return Return a vector of size `N` containing the decoded results.
*/

21
sherpa-onnx/csrc/offline-paraformer-greedy-search-decoder.cc
View File

@@ -4,28 +4,33 @@
#include "sherpa-onnx/csrc/offline-paraformer-greedy-search-decoder.h"
#include <algorithm>
#include <vector>
namespace sherpa_onnx {
std::vector<OfflineParaformerDecoderResult>
OfflineParaformerGreedySearchDecoder::Decode(Ort::Value /*log_probs*/,
Ort::Value token_num) {
std::vector<int64_t> shape = token_num.GetTensorTypeAndShapeInfo().GetShape();
OfflineParaformerGreedySearchDecoder::Decode(Ort::Value log_probs,
Ort::Value /*token_num*/) {
std::vector<int64_t> shape = log_probs.GetTensorTypeAndShapeInfo().GetShape();
int32_t batch_size = shape[0];
int32_t num_tokens = shape[1];
int32_t vocab_size = shape[2];
std::vector<OfflineParaformerDecoderResult> results(batch_size);
const int64_t *p = token_num.GetTensorData<int64_t>();
for (int32_t i = 0; i != batch_size; ++i) {
const float *p =
log_probs.GetTensorData<float>() + i * num_tokens * vocab_size;
for (int32_t k = 0; k != num_tokens; ++k) {
if (p[k] == eos_id_) break;
auto max_idx = static_cast<int64_t>(
std::distance(p, std::max_element(p, p + vocab_size)));
if (max_idx == eos_id_) break;
results[i].tokens.push_back(p[k]);
results[i].tokens.push_back(max_idx);
p += vocab_size;
}
p += num_tokens;
}
return results;

2
sherpa-onnx/csrc/offline-paraformer-greedy-search-decoder.h
View File

@@ -17,7 +17,7 @@ class OfflineParaformerGreedySearchDecoder : public OfflineParaformerDecoder {
: eos_id_(eos_id) {}
std::vector<OfflineParaformerDecoderResult> Decode(
Ort::Value /*log_probs*/, Ort::Value token_num) override;
Ort::Value log_probs, Ort::Value /*token_num*/) override;
private:
int32_t eos_id_;