diff --git a/src/config.cpp b/src/config.cpp index 4413d830c6..a9ec271835 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -1202,6 +1202,8 @@ struct Search_Element : JSON::Element { v_.past_present_share_buffer = JSON::Get(value); } else if (name == "early_stopping") { v_.early_stopping = JSON::Get(value); + } else if (name == "blank_penalty") { + v_.blank_penalty = static_cast(JSON::Get(value)); } else { throw JSON::unknown_value_error{}; } diff --git a/src/config.h b/src/config.h index 9d775c166c..f0bad1ccf0 100644 --- a/src/config.h +++ b/src/config.h @@ -379,6 +379,7 @@ struct Config { bool past_present_share_buffer{}; // The past/present kv tensors are shared and allocated once to max_length (cuda only) int random_seed{-1}; // -1 = Seed with random device, otherwise use value to seed RNG std::optional chunk_size; // Chunk size for prefill chunking during context processing. If present, chunking is enabled with the chunk size > 0. + float blank_penalty{}; // Penalty applied to blank token logits in CTC/RNNT decoding. Default 0 means no penalty. } search; struct Engine { diff --git a/src/models/nemotron_speech.cpp b/src/models/nemotron_speech.cpp index 8b560ec062..ee5fbba7bc 100644 --- a/src/models/nemotron_speech.cpp +++ b/src/models/nemotron_speech.cpp @@ -11,7 +11,7 @@ namespace Generators { -void NemotronCacheConfig::PopulateFromConfig(const Config& config) { +void NemotronConfig::PopulateFromConfig(const Config& config) { const auto& enc = config.model.encoder; const auto& dec = config.model.decoder; const auto& jo = config.model.joiner; @@ -39,6 +39,7 @@ void NemotronCacheConfig::PopulateFromConfig(const Config& config) { chunk_samples = config.model.chunk_samples; blank_id = config.model.blank_id; max_symbols_per_step = config.model.max_symbols_per_step; + blank_penalty = config.search.blank_penalty; // Vocab size from top-level config vocab_size = config.model.vocab_size; @@ -71,7 +72,7 @@ void NemotronCacheConfig::PopulateFromConfig(const Config& config) { dec_out_lstm_cell = dec.outputs.lstm_cell_state; } -void NemotronEncoderCache::Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { +void NemotronEncoderCache::Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { auto cache_channel_type = session_info.GetInputDataType(cfg.enc_in_cache_channel); auto cache_time_type = session_info.GetInputDataType(cfg.enc_in_cache_time); auto cache_channel_len_type = session_info.GetInputDataType(cfg.enc_in_cache_channel_len); @@ -92,11 +93,11 @@ void NemotronEncoderCache::Initialize(const NemotronCacheConfig& cfg, const Sess *cache_last_channel_len->GetTensorMutableData() = 0; } -void NemotronEncoderCache::Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { +void NemotronEncoderCache::Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { Initialize(cfg, session_info, allocator, device); } -void NemotronDecoderState::Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { +void NemotronDecoderState::Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { auto lstm_hidden_type = session_info.GetInputDataType(cfg.dec_in_lstm_hidden); auto lstm_cell_type = session_info.GetInputDataType(cfg.dec_in_lstm_cell); @@ -111,14 +112,14 @@ void NemotronDecoderState::Initialize(const NemotronCacheConfig& cfg, const Sess last_token = cfg.blank_id; // Start with blank/SOS token } -void NemotronDecoderState::Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { +void NemotronDecoderState::Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) { Initialize(cfg, session_info, allocator, device); } NemotronSpeechModel::NemotronSpeechModel(std::unique_ptr config, OrtEnv& ort_env) : Model{std::move(config)} { - cache_config_ = NemotronCacheConfig{}; - cache_config_.PopulateFromConfig(*config_); + nemotron_config_ = NemotronConfig{}; + nemotron_config_.PopulateFromConfig(*config_); // Create session options encoder_session_options_ = OrtSessionOptions::Create(); @@ -169,7 +170,7 @@ std::unique_ptr NemotronSpeechModel::CreateState(DeviceSpan /*se NemotronEncoderSubState::NemotronEncoderSubState(const NemotronSpeechModel& model, const GeneratorParams& params) : State{params, model}, model_{model} { - auto& cfg = model_.cache_config_; + auto& cfg = model_.nemotron_config_; auto& allocator = model_.allocator_cpu_; auto& device = *model_.p_device_; @@ -250,7 +251,7 @@ DeviceSpan NemotronEncoderSubState::Run(int /*total_length*/, DeviceSpan< NemotronPredictionSubState::NemotronPredictionSubState(const NemotronSpeechModel& model, const GeneratorParams& params) : State{params, model}, model_{model} { - auto& cfg = model_.cache_config_; + auto& cfg = model_.nemotron_config_; auto& allocator = model_.allocator_cpu_; auto& device = *model_.p_device_; @@ -304,7 +305,7 @@ DeviceSpan NemotronPredictionSubState::Run(int /*total_length*/, DeviceSp NemotronJoinerSubState::NemotronJoinerSubState(const NemotronSpeechModel& model, const GeneratorParams& params) : State{params, model}, model_{model} { - auto& cfg = model_.cache_config_; + auto& cfg = model_.nemotron_config_; // Register inputs encoder_input_idx_ = inputs_.size(); @@ -339,15 +340,15 @@ NemotronSpeechState::NemotronSpeechState(const NemotronSpeechModel& model, const GeneratorParams& params) : State{params, model}, nemotron_model_{model} { - cache_config_ = model.cache_config_; + nemotron_config_ = model.nemotron_config_; encoder_state_ = std::make_unique(model, params); prediction_state_ = std::make_unique(model, params); joiner_state_ = std::make_unique(model, params); // Pre-allocate encoder frame for joiner input - auto enc_out_type = model_.session_info_.GetOutputDataType(cache_config_.enc_out_encoded); - auto frame_shape = std::array{1, 1, cache_config_.hidden_dim}; + auto enc_out_type = model_.session_info_.GetOutputDataType(nemotron_config_.enc_out_encoded); + auto frame_shape = std::array{1, 1, nemotron_config_.hidden_dim}; encoder_frame_ = OrtValue::CreateTensor(model_.allocator_cpu_, frame_shape, enc_out_type); } @@ -363,7 +364,7 @@ DeviceSpan NemotronSpeechState::Run(int /*total_length*/, void NemotronSpeechState::SetExtraInputs(const std::vector& extra_inputs) { for (const auto& input : extra_inputs) { - if (input.name == Config::Defaults::AudioFeaturesName || input.name == cache_config_.enc_in_audio) { + if (input.name == Config::Defaults::AudioFeaturesName || input.name == nemotron_config_.enc_in_audio) { current_mel_ = input.tensor; need_encoder_run_ = true; chunk_done_ = false; @@ -389,11 +390,11 @@ void NemotronSpeechState::ResetStreamingState() { auto& allocator = model_.allocator_cpu_; auto& device = *model_.p_device_; - encoder_state_->cache_.Reset(cache_config_, model_.session_info_, allocator, device); + encoder_state_->cache_.Reset(nemotron_config_, model_.session_info_, allocator, device); encoder_state_->UpdateCacheInputs(); encoder_state_->first_run_ = true; - prediction_state_->lstm_state_.Reset(cache_config_, model_.session_info_, allocator, device); + prediction_state_->lstm_state_.Reset(nemotron_config_, model_.session_info_, allocator, device); prediction_state_->UpdateInputs(); prediction_state_->first_run_ = true; @@ -473,7 +474,7 @@ std::span NemotronSpeechState::StepToken() { // Reshape decoder output for joiner: [1, dim] -> [1, 1, dim] auto dec_out_shape = prediction_state_->outputs_[0]->GetTensorTypeAndShapeInfo()->GetShape(); auto decoder_frame_shape = std::array{1, 1, dec_out_shape[1]}; - auto dec_out_type = model_.session_info_.GetOutputDataType(cache_config_.dec_out_outputs); + auto dec_out_type = model_.session_info_.GetOutputDataType(nemotron_config_.dec_out_outputs); auto decoder_frame = OrtValue::CreateTensor(allocator, decoder_frame_shape, dec_out_type); ByteWrapTensor(*model_.p_device_, *decoder_frame) .CopyFrom(ByteWrapTensor(*model_.p_device_, *prediction_state_->outputs_[0])); @@ -483,21 +484,23 @@ std::span NemotronSpeechState::StepToken() { joiner_state_->Run(0, dummy_tokens); // Argmax over logits - const float* logits_data = joiner_state_->outputs_[0]->GetTensorData(); + const float* logits = joiner_state_->outputs_[0]->GetTensorData(); auto logits_shape = joiner_state_->outputs_[0]->GetTensorTypeAndShapeInfo()->GetShape(); int total_logits = 1; for (auto d : logits_shape) total_logits *= static_cast(d); + // Apply blank penalty virtually during argmax to avoid mutating ORT output buffer int best_token = 0; - float best_score = logits_data[0]; + float best_score = logits[0] - (nemotron_config_.blank_id == 0 ? nemotron_config_.blank_penalty : 0.0f); for (int i = 1; i < total_logits; ++i) { - if (logits_data[i] > best_score) { - best_score = logits_data[i]; + float score = (i == nemotron_config_.blank_id) ? logits[i] - nemotron_config_.blank_penalty : logits[i]; + if (score > best_score) { + best_score = score; best_token = i; } } - if (best_token == cache_config_.blank_id) { + if (best_token == nemotron_config_.blank_id) { time_step_++; symbol_step_ = 0; continue; @@ -511,7 +514,7 @@ std::span NemotronSpeechState::StepToken() { prediction_state_->outputs_[2] = nullptr; symbol_step_++; - if (symbol_step_ >= cache_config_.max_symbols_per_step) { + if (symbol_step_ >= nemotron_config_.max_symbols_per_step) { time_step_++; symbol_step_ = 0; } diff --git a/src/models/nemotron_speech.h b/src/models/nemotron_speech.h index 43ddd75021..8b58454586 100644 --- a/src/models/nemotron_speech.h +++ b/src/models/nemotron_speech.h @@ -9,7 +9,7 @@ namespace Generators { -struct NemotronCacheConfig { +struct NemotronConfig { // Encoder dimensions (from encoder.hidden_size / num_hidden_layers) int num_encoder_layers{}; int hidden_dim{}; @@ -23,6 +23,7 @@ struct NemotronCacheConfig { // Vocabulary int vocab_size{}; int blank_id{}; + float blank_penalty{}; // Streaming chunk config int chunk_frames{}; @@ -76,8 +77,8 @@ struct NemotronEncoderCache { std::unique_ptr cache_last_time; std::unique_ptr cache_last_channel_len; - void Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); - void Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); + void Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); + void Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); }; /// Holds the RNNT decoder LSTM hidden states between decoding steps. @@ -86,8 +87,8 @@ struct NemotronDecoderState { std::unique_ptr lstm_cell_state; int last_token{0}; // Last emitted non-blank token (for autoregressive feedback) - void Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); - void Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); + void Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); + void Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device); }; struct NemotronSpeechModel : Model { @@ -105,7 +106,7 @@ struct NemotronSpeechModel : Model { std::unique_ptr decoder_session_options_; std::unique_ptr joiner_session_options_; - NemotronCacheConfig cache_config_; + NemotronConfig nemotron_config_; }; /// Sub-state for the streaming encoder. @@ -201,7 +202,7 @@ struct NemotronSpeechState : State { private: const NemotronSpeechModel& nemotron_model_; - NemotronCacheConfig cache_config_; + NemotronConfig nemotron_config_; std::unique_ptr encoder_state_; std::unique_ptr prediction_state_; diff --git a/src/models/nemotron_streaming_processor.cpp b/src/models/nemotron_streaming_processor.cpp index f9a5cbfe90..a0284ff355 100644 --- a/src/models/nemotron_streaming_processor.cpp +++ b/src/models/nemotron_streaming_processor.cpp @@ -16,24 +16,24 @@ NemotronStreamingProcessor::NemotronStreamingProcessor(Model& model) throw std::runtime_error("NemotronStreamingProcessor requires a nemotron_speech model type. Got: " + model.config_->model.type); } - cache_config_ = nemotron_model->cache_config_; + nemotron_config_ = nemotron_model->nemotron_config_; - if (cache_config_.pre_encode_cache_size <= 0) { + if (nemotron_config_.pre_encode_cache_size <= 0) { throw std::runtime_error("NemotronStreamingProcessor requires pre_encode_cache_size > 0. Got: " + - std::to_string(cache_config_.pre_encode_cache_size)); + std::to_string(nemotron_config_.pre_encode_cache_size)); } // Initialize mel extractor from config nemo_mel::NemoMelConfig mel_cfg{ - cache_config_.num_mels, cache_config_.fft_size, - cache_config_.hop_length, cache_config_.win_length, - cache_config_.sample_rate, - cache_config_.preemph, cache_config_.log_eps}; + nemotron_config_.num_mels, nemotron_config_.fft_size, + nemotron_config_.hop_length, nemotron_config_.win_length, + nemotron_config_.sample_rate, + nemotron_config_.preemph, nemotron_config_.log_eps}; mel_extractor_ = nemo_mel::NemoStreamingMelExtractor{mel_cfg}; // Initialize mel pre-encode cache (time-major ring buffer, zeros for first chunk) mel_pre_encode_cache_.assign( - static_cast(cache_config_.pre_encode_cache_size) * cache_config_.num_mels, 0.0f); + static_cast(nemotron_config_.pre_encode_cache_size) * nemotron_config_.num_mels, 0.0f); cache_pos_ = 0; } @@ -43,7 +43,7 @@ std::unique_ptr NemotronStreamingProcessor::Process(const float* a // Append incoming audio to accumulation buffer audio_buffer_.insert(audio_buffer_.end(), audio_data, audio_data + num_samples); - const size_t chunk_size = static_cast(cache_config_.chunk_samples); + const size_t chunk_size = static_cast(nemotron_config_.chunk_samples); // Process the first complete chunk available if (audio_buffer_.size() >= chunk_size) { @@ -63,7 +63,7 @@ std::unique_ptr NemotronStreamingProcessor::Flush() { return nullptr; } - const size_t chunk_size = static_cast(cache_config_.chunk_samples); + const size_t chunk_size = static_cast(nemotron_config_.chunk_samples); audio_buffer_.resize(chunk_size, 0.0f); // Pad with silence auto mel = BuildMelTensor(audio_buffer_.data(), chunk_size); @@ -79,12 +79,12 @@ std::unique_ptr NemotronStreamingProcessor::BuildMelTensor(const float // Compute mel spectrogram for this chunk: returns [num_mels, num_frames] (frequency-major) auto [mel_data, num_frames] = mel_extractor_.Process(audio_chunk, chunk_samples); - const int cache_size = cache_config_.pre_encode_cache_size; - const int num_mels = cache_config_.num_mels; + const int cache_size = nemotron_config_.pre_encode_cache_size; + const int num_mels = nemotron_config_.num_mels; const int total_mel_frames = cache_size + num_frames; // Create output tensor: [1, total_mel_frames, num_mels] (time-major) - auto signal_type = model_.session_info_.GetInputDataType(cache_config_.enc_in_audio); + auto signal_type = model_.session_info_.GetInputDataType(nemotron_config_.enc_in_audio); // TODO: Optimize for GPU/CUDA later, CPU always expects float32. if (signal_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { diff --git a/src/models/nemotron_streaming_processor.h b/src/models/nemotron_streaming_processor.h index 405f2b79c0..27da5f2b75 100644 --- a/src/models/nemotron_streaming_processor.h +++ b/src/models/nemotron_streaming_processor.h @@ -19,12 +19,12 @@ struct NemotronStreamingProcessor : StreamingProcessor { std::unique_ptr Process(const float* audio_data, size_t num_samples) override; std::unique_ptr Flush() override; - int GetChunkSamples() const { return cache_config_.chunk_samples; } - int GetSampleRate() const { return cache_config_.sample_rate; } + int GetChunkSamples() const { return nemotron_config_.chunk_samples; } + int GetSampleRate() const { return nemotron_config_.sample_rate; } private: Model& model_; - NemotronCacheConfig cache_config_; + NemotronConfig nemotron_config_; // Log-mel feature extraction nemo_mel::NemoStreamingMelExtractor mel_extractor_;