Skip to content

Fix bugs in NeMo transducer modified beam search. #3589

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
csukuangfj merged 1 commit into from
May 8, 2026
Merged

Fix bugs in NeMo transducer modified beam search. #3589

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 31 additions & 2 deletions sherpa-onnx/csrc/offline-transducer-modified-beam-search-nemo-decoder.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,14 @@ struct NeMoHypothesis {
const ContextState *context_state; // context graph state
OrtAllocator *allocator; // allocator for cloning states
int32_t frame_offset; // current frame position for this hypothesis
int32_t num_symbols; // number of non-blank symbols emitted at current frame

NeMoHypothesis()
: log_prob(0.0f),
context_state(nullptr),
allocator(nullptr),
frame_offset(0) {}
frame_offset(0),
num_symbols(0) {}

// Copy constructor - needed for hypothesis expansion
NeMoHypothesis(const NeMoHypothesis &other)
Expand All @@ -45,7 +47,8 @@ struct NeMoHypothesis {
log_prob(other.log_prob),
context_state(other.context_state),
allocator(other.allocator),
frame_offset(other.frame_offset) {
frame_offset(other.frame_offset),
num_symbols(other.num_symbols) {
// Deep copy of decoder states
decoder_states.reserve(other.decoder_states.size());
for (const auto &state : other.decoder_states) {
Expand All @@ -63,6 +66,7 @@ struct NeMoHypothesis {
context_state = other.context_state;
allocator = other.allocator;
frame_offset = other.frame_offset;
num_symbols = other.num_symbols;

decoder_states.clear();
decoder_states.reserve(other.decoder_states.size());
Expand Down Expand Up @@ -90,6 +94,7 @@ OfflineTransducerModifiedBeamSearchNeMoDecoder::Decode(

int32_t vocab_size = model_->VocabSize();
int32_t blank_id = vocab_size - 1; // NeMo models have blank at the end
int32_t max_symbols_per_frame = 10;

@gemini-code-assist gemini-code-assist Bot May 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The max_symbols_per_frame value is hardcoded. It would be better to make this a configurable parameter in the decoder options to allow users to tune the behavior for different models. At the very least, it should be marked as constexpr to indicate it is a constant.

Suggested change
int32_t max_symbols_per_frame = 10;
constexpr int32_t kMaxSymbolsPerFrame = 10;


// For TDT models, we need to know the number of duration bins
// We'll detect this from the joiner output size on first run
Expand Down Expand Up @@ -168,6 +173,26 @@ OfflineTransducerModifiedBeamSearchNeMoDecoder::Decode(
continue;
}

if (hyp.num_symbols >= max_symbols_per_frame) {
// Reached the per-frame symbol limit, force blank to advance frame
NeMoHypothesis new_hyp;
new_hyp.ys = hyp.ys;
new_hyp.timestamps = hyp.timestamps;
new_hyp.durations = hyp.durations;
new_hyp.ys_probs = hyp.ys_probs;
new_hyp.context_state = hyp.context_state;
new_hyp.allocator = allocator;
new_hyp.log_prob = hyp.log_prob;
new_hyp.num_symbols = 0;
new_hyp.decoder_states.reserve(hyp.decoder_states.size());
for (const auto &state : hyp.decoder_states) {
new_hyp.decoder_states.push_back(Clone(allocator, &state));
}
new_hyp.frame_offset = hyp.frame_offset + 1;
all_candidates.emplace_back(new_hyp.log_prob, std::move(new_hyp));
continue;
Comment on lines +178 to +193

@gemini-code-assist gemini-code-assist Bot May 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The manual copying of NeMoHypothesis fields and deep copying of decoder_states is redundant and inefficient. Since hyp is not used again in this iteration after the continue, you can use std::move(hyp) to avoid expensive cloning of ONNX tensors. This significantly improves performance as it avoids memory allocations and data copying for the decoder states.

Suggested change
NeMoHypothesis new_hyp;
new_hyp.ys = hyp.ys;
new_hyp.timestamps = hyp.timestamps;
new_hyp.durations = hyp.durations;
new_hyp.ys_probs = hyp.ys_probs;
new_hyp.context_state = hyp.context_state;
new_hyp.allocator = allocator;
new_hyp.log_prob = hyp.log_prob;
new_hyp.num_symbols = 0;
new_hyp.decoder_states.reserve(hyp.decoder_states.size());
for (const auto &state : hyp.decoder_states) {
new_hyp.decoder_states.push_back(Clone(allocator, &state));
}
new_hyp.frame_offset = hyp.frame_offset + 1;
all_candidates.emplace_back(new_hyp.log_prob, std::move(new_hyp));
continue;
NeMoHypothesis new_hyp = std::move(hyp);
new_hyp.num_symbols = 0;
new_hyp.frame_offset += 1;
all_candidates.emplace_back(new_hyp.log_prob, std::move(new_hyp));
continue;

}

// Get encoder output for this frame
std::array<int64_t, 3> encoder_3d_shape{1, encoder_dim, 1};
const float *frame_data = this_encoder + hyp.frame_offset * encoder_dim;
Expand Down Expand Up @@ -295,6 +320,7 @@ OfflineTransducerModifiedBeamSearchNeMoDecoder::Decode(
// For blank/unk in TDT, always advance by at least 1
new_hyp.frame_offset =
hyp.frame_offset + std::max(1, predicted_skip);
new_hyp.num_symbols = 0;
} else {
// Non-blank: add token, use new decoder state
new_hyp.ys.push_back(token);
Expand All @@ -314,8 +340,11 @@ OfflineTransducerModifiedBeamSearchNeMoDecoder::Decode(
// tokens
if (is_tdt_) {
new_hyp.frame_offset = hyp.frame_offset + predicted_skip;
new_hyp.num_symbols =
predicted_skip > 0 ? 0 : hyp.num_symbols + 1;
} else {
new_hyp.frame_offset = hyp.frame_offset;
new_hyp.num_symbols = hyp.num_symbols + 1;
}

// Update context graph
Expand Down
2 changes: 1 addition & 1 deletion sherpa-onnx/csrc/online-recognizer-paraformer-impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ class OnlineRecognizerParaformerImpl : public OnlineRecognizerImpl {
OnlineRecognizerParaformerImpl(const OnlineRecognizerParaformerImpl &) =
delete;

OnlineRecognizerParaformerImpl operator=(
OnlineRecognizerParaformerImpl &operator=(
const OnlineRecognizerParaformerImpl &) = delete;

std::unique_ptr<OnlineStream> CreateStream() const override {
Expand Down
Loading