Skip to content

Commit

Permalink
Showing 3 changed files with 15 additions and 11 deletions.
3 changes: 1 addition & 2 deletions sherpa-onnx/csrc/offline-paraformer-decoder.h
Original file line number Diff line number Diff line change
@@ -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 changes: 13 additions & 8 deletions sherpa-onnx/csrc/offline-paraformer-greedy-search-decoder.cc
Original file line number Diff line number Diff line change
@@ -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 += num_tokens;
p += vocab_size;
}
}

return results;
Original file line number Diff line number Diff line change
@@ -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_;

0 comments on commit d781fcd

Please sign in to comment.