From 1892ba7c2b23af0a04372bc8e195dca17f6e23e3 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Thu, 7 May 2026 21:16:38 +0000 Subject: [PATCH 01/34] Start parakeet v3 --- examples/python/parakeet_simulate_mic.py | 116 +++++ src/config.cpp | 25 ++ src/config.h | 16 + src/generators.cpp | 1 + src/leakcheck.h | 3 +- src/models/model.cpp | 3 + src/models/model_type.h | 6 + src/models/parakeet_speech.cpp | 150 +++++++ src/models/parakeet_speech.h | 117 +++++ src/ort_genai.h | 32 ++ src/ort_genai_c.cpp | 42 ++ src/ort_genai_c.h | 34 ++ src/parakeet_mel.cpp | 249 +++++++++++ src/parakeet_mel.h | 42 ++ src/parakeet_streaming_asr.cpp | 532 +++++++++++++++++++++++ src/parakeet_streaming_asr.h | 72 +++ src/python/python.cpp | 21 + src/streaming_asr.h | 45 ++ 18 files changed, 1505 insertions(+), 1 deletion(-) create mode 100644 examples/python/parakeet_simulate_mic.py create mode 100644 src/models/parakeet_speech.cpp create mode 100644 src/models/parakeet_speech.h create mode 100644 src/parakeet_mel.cpp create mode 100644 src/parakeet_mel.h create mode 100644 src/parakeet_streaming_asr.cpp create mode 100644 src/parakeet_streaming_asr.h create mode 100644 src/streaming_asr.h diff --git a/examples/python/parakeet_simulate_mic.py b/examples/python/parakeet_simulate_mic.py new file mode 100644 index 0000000000..c74476c7b7 --- /dev/null +++ b/examples/python/parakeet_simulate_mic.py @@ -0,0 +1,116 @@ +""" +Parakeet TDT Speech Streaming ASR — Simulated real-time microphone demo. + +Reads an audio file and feeds it chunk-by-chunk in real-time, +simulating live microphone input with actual wall-clock delays. + +Usage: + python parakeet_simulate_mic.py --model_path ./parakeet-tdt-0.6b-v3-onnx-fp32/fp32 --audio_file recording.wav +""" + +import argparse +import os +import sys +import time +import numpy as np + +import onnxruntime_genai as og + +SAMPLE_RATE = 16000 +CHUNK_SAMPLES = 12800 # 0.8 seconds at 16kHz (NeMo recommended default) +CHUNK_DURATION = CHUNK_SAMPLES / SAMPLE_RATE + + +def load_audio(audio_path): + """Load a WAV file as float32 mono at 16kHz.""" + try: + import soundfile as sf + except ImportError: + raise ImportError("pip install soundfile") + + audio, sr = sf.read(audio_path, dtype="float32") + + if len(audio.shape) > 1: + audio = audio.mean(axis=1) + + if sr != SAMPLE_RATE: + try: + import scipy.signal + num_samples = int(len(audio) * SAMPLE_RATE / sr) + audio = scipy.signal.resample(audio, num_samples).astype(np.float32) + print(f" Resampled {sr}Hz -> {SAMPLE_RATE}Hz") + except ImportError: + raise ValueError(f"Expected {SAMPLE_RATE}Hz, got {sr}Hz. pip install scipy") + + return audio + + +def simulate_microphone(model_path, audio_path, realtime=True): + """Simulate real-time microphone streaming with wall-clock delays.""" + + print("=" * 60) + print(" PARAKEET TDT — SIMULATED REAL-TIME MICROPHONE") + print("=" * 60) + + # Load audio + print(f"\nLoading audio: {audio_path}") + audio = load_audio(audio_path) + duration = len(audio) / SAMPLE_RATE + num_chunks = (len(audio) + CHUNK_SAMPLES - 1) // CHUNK_SAMPLES + print(f" Duration: {duration:.1f}s | Chunks: {num_chunks} x {CHUNK_DURATION*1000:.0f}ms") + + # Load model + print(f"\nLoading model: {model_path}") + config = og.Config(model_path) + model = og.Model(config) + asr = og.StreamingASR(model) + print(" Model ready.\n") + + # Simulate streaming + print("-" * 60) + print("LIVE TRANSCRIPTION (simulated real-time):") + print("-" * 60) + print() + + stream_start = time.time() + + for i in range(0, len(audio), CHUNK_SAMPLES): + # Wait for real-time moment (simulate microphone latency) + if realtime: + audio_time = i / SAMPLE_RATE + target_time = stream_start + audio_time + now = time.time() + + # Feed chunk + chunk = audio[i:i + CHUNK_SAMPLES] + if len(chunk) < CHUNK_SAMPLES: + chunk = np.pad(chunk, (0, CHUNK_SAMPLES - len(chunk))) + chunk = chunk.astype(np.float32) + + raw_text = asr.transcribe_chunk(chunk) + + if raw_text: + print(raw_text, end="", flush=True) + + # Flush remaining audio + flush_text = asr.flush() + if flush_text: + print(flush_text, end="", flush=True) + + elapsed = time.time() - stream_start + print(f"\n{'='*60}") + print(f"Audio duration: {duration:.1f}s | Processing: {elapsed:.1f}s | RTF: {elapsed/duration:.2f}x") + print(f"\nFull transcript:") + print(asr.get_transcript()) + print(f"{'='*60}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Parakeet TDT Speech Streaming ASR") + parser.add_argument("--model_path", required=True, help="Path to model directory") + parser.add_argument("--audio_file", required=True, help="Path to audio file (WAV)") + parser.add_argument("--no-realtime", action="store_true", + help="Skip wall-clock delays (process as fast as possible)") + args = parser.parse_args() + + simulate_microphone(args.model_path, args.audio_file, realtime=not args.no_realtime) diff --git a/src/config.cpp b/src/config.cpp index 7778e9bf4e..1df8b90366 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -340,6 +340,12 @@ struct DecoderInputs_Element : JSON::Element { v_.lstm_hidden_state = JSON::Get(value); } else if (name == "lstm_cell_state") { v_.lstm_cell_state = JSON::Get(value); + } else if (name == "target_length") { + v_.target_length = JSON::Get(value); + } else if (name == "states_1") { + v_.states_1 = JSON::Get(value); + } else if (name == "states_2") { + v_.states_2 = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } @@ -373,6 +379,12 @@ struct DecoderOutputs_Element : JSON::Element { v_.lstm_hidden_state = JSON::Get(value); } else if (name == "lstm_cell_state") { v_.lstm_cell_state = JSON::Get(value); + } else if (name == "prednet_lengths") { + v_.prednet_lengths = JSON::Get(value); + } else if (name == "states_1") { + v_.states_1 = JSON::Get(value); + } else if (name == "states_2") { + v_.states_2 = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } @@ -1145,6 +1157,16 @@ struct Model_Element : JSON::Element { v_.blank_id = static_cast(JSON::Get(value)); } else if (name == "max_symbols_per_step") { v_.max_symbols_per_step = static_cast(JSON::Get(value)); + } else if (name == "left_context_samples") { + v_.left_context_samples = static_cast(JSON::Get(value)); + } else if (name == "right_context_samples") { + v_.right_context_samples = static_cast(JSON::Get(value)); + } else if (name == "tdt_num_extra_outputs") { + v_.tdt_num_extra_outputs = static_cast(JSON::Get(value)); + } else if (name == "enc_in_length") { + v_.enc_in_length = JSON::Get(value); + } else if (name == "enc_out_length") { + v_.enc_out_length = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } @@ -1153,6 +1175,8 @@ struct Model_Element : JSON::Element { Element& OnArray(std::string_view name) override { if (name == "eos_token_id") return eos_token_id_; + if (name == "tdt_durations") + return tdt_durations_; throw JSON::unknown_value_error{}; } @@ -1186,6 +1210,7 @@ struct Model_Element : JSON::Element { Encoder_Element encoder_{v_.encoder}; Decoder_Element decoder_{v_.decoder}; Int_Array_Element eos_token_id_{v_.eos_token_id}; + Int_Array_Element tdt_durations_{v_.tdt_durations}; Vision_Element vision_{v_.vision}; Embedding_Element embedding_{v_.embedding}; Speech_Element speech_{v_.speech}; diff --git a/src/config.h b/src/config.h index 155970ba22..a130beef0b 100644 --- a/src/config.h +++ b/src/config.h @@ -154,6 +154,14 @@ struct Config { int blank_id{}; int max_symbols_per_step{}; + // Parakeet TDT (Token-and-Duration Transducer) parameters + int left_context_samples{}; // Left context in PCM samples (e.g., 143360 for 9.0s) + int right_context_samples{}; // Right context in PCM samples (e.g., 25600 for 1.6s) + std::vector tdt_durations; // e.g., {0, 1, 2, 3, 4} + int tdt_num_extra_outputs{}; // Number of duration logit outputs (e.g., 5) + std::string enc_in_length{"length"}; + std::string enc_out_length{"encoded_lengths"}; + struct Encoder { std::string filename; std::optional session_options; @@ -346,6 +354,10 @@ struct Config { std::string targets; std::string lstm_hidden_state; std::string lstm_cell_state; + // Parakeet TDT decoder (prediction network) extra inputs + std::string target_length; // "target_length" + std::string states_1; // "states.1" + std::string states_2; // "onnx::Slice_3" } inputs; struct Outputs { @@ -361,6 +373,10 @@ struct Config { std::string outputs; std::string lstm_hidden_state; std::string lstm_cell_state; + // Parakeet TDT decoder (prediction network) extra outputs + std::string prednet_lengths; // "prednet_lengths" + std::string states_1; // "states" + std::string states_2; // "162" } outputs; struct PipelineModel { diff --git a/src/generators.cpp b/src/generators.cpp index 65d1bef8e6..955c38135b 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -4,6 +4,7 @@ // Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved. #include "generators.h" +#include "streaming_asr.h" #include "models/streaming_processor.h" #include "models/nemotron_speech.h" #include "sequences.h" diff --git a/src/leakcheck.h b/src/leakcheck.h index 8be2652454..0832454bd9 100644 --- a/src/leakcheck.h +++ b/src/leakcheck.h @@ -14,6 +14,7 @@ struct Generator; struct Model; struct Request; struct Search; +struct StreamingASR; struct StreamingProcessor; struct Tensor; struct Tokenizer; @@ -26,7 +27,7 @@ struct LeakTypeList { static bool Dump(); }; -using LeakTypes = LeakTypeList; +using LeakTypes = LeakTypeList; template struct LeakChecked { diff --git a/src/models/model.cpp b/src/models/model.cpp index 8a4c248cf5..e3134895d7 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -16,6 +16,7 @@ #include "gpt.h" #include "decoder_only.h" #include "whisper.h" +#include "parakeet_speech.h" #include "nemotron_speech.h" #include "multi_modal.h" #include "lfm2.h" @@ -830,6 +831,8 @@ std::shared_ptr CreateModel(OrtEnv& ort_env, std::unique_ptr conf return std::make_shared(std::move(config), ort_env); if (ModelType::IsRNNT(config->model.type)) return std::make_shared(std::move(config), ort_env); + if (ModelType::IsStreamingASR(config->model.type)) + return std::make_shared(std::move(config), ort_env); if (ModelType::IsALM(config->model.type)) return std::make_shared(std::move(config), ort_env); if (ModelType::IsVLM(config->model.type)) diff --git a/src/models/model_type.h b/src/models/model_type.h index e34dd348f0..21315492c9 100644 --- a/src/models/model_type.h +++ b/src/models/model_type.h @@ -47,6 +47,12 @@ struct ModelType { return std::find(rnnt_types.begin(), rnnt_types.end(), model_type) != rnnt_types.end(); } + inline static bool IsStreamingASR(const std::string& model_type) { + // Streaming ASR model (cache-aware encoder + RNNT/TDT decoder) + static constexpr std::array StreamingASR = {"parakeet_tdt"}; + return std::find(StreamingASR.begin(), StreamingASR.end(), model_type) != StreamingASR.end(); + } + inline static bool IsMMM(const std::string& model_type) { // Multi-modal model (MMM) static constexpr std::array MMM = {"gemma4", "phi4mm"}; diff --git a/src/models/parakeet_speech.cpp b/src/models/parakeet_speech.cpp new file mode 100644 index 0000000000..5ada40be6d --- /dev/null +++ b/src/models/parakeet_speech.cpp @@ -0,0 +1,150 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Parakeet TDT Speech Model — non-cache-aware encoder + TDT decoder+joiner. + +#include +#include +#include + +#include "../generators.h" +#include "parakeet_speech.h" + +namespace Generators { + +// ─── ParakeetConfig ───────────────────────────────────────────────────────── + +void ParakeetConfig::PopulateFromConfig(const Config& config) { + const auto& enc = config.model.encoder; + const auto& dec = config.model.decoder; + const auto& m = config.model; + const auto& jo = config.model.joiner; + + // Encoder dimensions + hidden_dim = enc.hidden_size; + num_encoder_layers = enc.num_hidden_layers; + + // Decoder dimensions (LSTM) + decoder_lstm_dim = dec.hidden_size; + decoder_lstm_layers = dec.num_hidden_layers; + + // Speech / mel feature config (top-level model.* in this repo) + num_mels = m.num_mels; + fft_size = m.fft_size; + hop_length = m.hop_length; + win_length = m.win_length; + preemph = m.preemph; + log_eps = m.log_eps; + subsampling_factor = m.subsampling_factor; + sample_rate = m.sample_rate; + chunk_samples = m.chunk_samples; + blank_id = m.blank_id; + max_symbols_per_step = m.max_symbols_per_step; + left_context_samples = m.left_context_samples; + right_context_samples = m.right_context_samples; + + // TDT parameters + tdt_durations = m.tdt_durations; + tdt_num_extra_outputs = m.tdt_num_extra_outputs; + + // Vocab size from top-level config + vocab_size = config.model.vocab_size; + + // Encoder I/O names + enc_in_audio = enc.inputs.audio_features; + enc_out_encoded = enc.outputs.encoder_outputs; + enc_in_length = m.enc_in_length; + enc_out_length = m.enc_out_length; + + // Joiner I/O names + join_in_encoder = jo.inputs.encoder_outputs; + join_in_decoder = jo.inputs.decoder_outputs; + join_out_logits = jo.outputs.logits; + + // Decoder I/O names (RNNT/TDT prediction network) + dec_in_targets = dec.inputs.targets; + dec_in_target_length = dec.inputs.target_length; + dec_in_states_1 = dec.inputs.states_1; + dec_in_states_2 = dec.inputs.states_2; + dec_out_outputs = dec.outputs.outputs; + dec_out_prednet_lengths = dec.outputs.prednet_lengths; + dec_out_states_1 = dec.outputs.states_1; + dec_out_states_2 = dec.outputs.states_2; +} + +// ─── ParakeetDecoderState ─────────────────────────────────────────────────── + +void ParakeetDecoderState::Initialize(const ParakeetConfig& cfg, OrtAllocator& allocator) { + // LSTM states: [lstm_layers, 1, lstm_dim] + auto state_shape = std::array{cfg.decoder_lstm_layers, 1, cfg.decoder_lstm_dim}; + state_h = OrtValue::CreateTensor(allocator, state_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + std::memset(state_h->GetTensorMutableRawData(), 0, + cfg.decoder_lstm_layers * 1 * cfg.decoder_lstm_dim * sizeof(float)); + + state_c = OrtValue::CreateTensor(allocator, state_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + std::memset(state_c->GetTensorMutableRawData(), 0, + cfg.decoder_lstm_layers * 1 * cfg.decoder_lstm_dim * sizeof(float)); + + last_token = static_cast(cfg.blank_id); + + // decoder_output will be initialized on first use via RunInitialDecoder + decoder_output.reset(); +} + +void ParakeetDecoderState::Reset(const ParakeetConfig& cfg, OrtAllocator& allocator) { + Initialize(cfg, allocator); +} + +// ─── ParakeetSpeechModel ──────────────────────────────────────────────────── + +ParakeetSpeechModel::ParakeetSpeechModel(std::unique_ptr config, OrtEnv& ort_env) + : Model{std::move(config)} { + // Populate from genai_config.json + parakeet_config_ = ParakeetConfig{}; + parakeet_config_.PopulateFromConfig(*config_); + + // Create session options + encoder_session_options_ = OrtSessionOptions::Create(); + decoder_session_options_ = OrtSessionOptions::Create(); + joiner_session_options_ = OrtSessionOptions::Create(); + + if (config_->model.encoder.session_options.has_value()) { + CreateSessionOptionsFromConfig(config_->model.encoder.session_options.value(), + *encoder_session_options_, true, false); + } else { + CreateSessionOptionsFromConfig(config_->model.decoder.session_options, + *encoder_session_options_, true, false); + } + CreateSessionOptionsFromConfig(config_->model.decoder.session_options, + *decoder_session_options_, true, false); + if (config_->model.joiner.session_options.has_value()) { + CreateSessionOptionsFromConfig(config_->model.joiner.session_options.value(), + *joiner_session_options_, true, false); + } else { + CreateSessionOptionsFromConfig(config_->model.decoder.session_options, + *joiner_session_options_, true, false); + } + + // Load the three ONNX models + std::string encoder_filename = config_->model.encoder.filename; + if (encoder_filename.empty()) encoder_filename = "encoder.onnx"; + + std::string decoder_filename = config_->model.decoder.filename; + if (decoder_filename.empty()) decoder_filename = "decoder.onnx"; + + std::string joiner_filename = config_->model.joiner.filename; + if (joiner_filename.empty()) joiner_filename = "joint.onnx"; + + session_encoder_ = CreateSession(ort_env, encoder_filename, encoder_session_options_.get()); + session_decoder_ = CreateSession(ort_env, decoder_filename, decoder_session_options_.get()); + session_joiner_ = CreateSession(ort_env, joiner_filename, joiner_session_options_.get()); +} + +std::unique_ptr ParakeetSpeechModel::CreateState(DeviceSpan /*sequence_lengths*/, + const GeneratorParams& /*params*/) const { + throw std::runtime_error( + "ParakeetSpeechModel does not support the Generator pipeline. " + "Use the StreamingASR API instead."); +} + +} // namespace Generators diff --git a/src/models/parakeet_speech.h b/src/models/parakeet_speech.h new file mode 100644 index 0000000000..1b181b18f4 --- /dev/null +++ b/src/models/parakeet_speech.h @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Parakeet TDT Speech Model support. +// Non-cache-aware encoder + TDT (Token-and-Duration Transducer) decoder+joiner +// for real-time streaming transcription. +#pragma once + +#include +#include +#include + +#include "model.h" + +namespace Generators { + +/// Configuration for Parakeet TDT model. +/// Populated from Config::Model at model load time via PopulateFromConfig(). +struct ParakeetConfig { + // Encoder dimensions + int num_encoder_layers{}; + int hidden_dim{}; + + // Decoder LSTM dimensions + int decoder_lstm_dim{}; + int decoder_lstm_layers{}; + + // Vocabulary + int vocab_size{}; + int blank_id{}; + + // Streaming chunk config + int sample_rate{16000}; + int chunk_samples{12800}; // 0.8 seconds at 16kHz (NeMo recommended default) + int subsampling_factor{8}; + int max_symbols_per_step{10}; + int left_context_samples{143360}; // ~9.0 seconds of left context + int right_context_samples{25600}; // 1.6 seconds of right context + + // Mel spectrogram parameters + int num_mels{}; + int fft_size{}; + int hop_length{}; + int win_length{}; + float preemph{}; + float log_eps{}; + + // TDT (Token-and-Duration Transducer) parameters + std::vector tdt_durations; // e.g., {0, 1, 2, 3, 4} + int tdt_num_extra_outputs{}; // Number of duration logit outputs (e.g., 5) + + // Encoder I/O names + std::string enc_in_audio; + std::string enc_in_length; + std::string enc_out_encoded; + std::string enc_out_length; + + // Decoder (prediction network) I/O names + std::string dec_in_targets; + std::string dec_in_target_length; + std::string dec_in_states_1; // h_in + std::string dec_in_states_2; // c_in + std::string dec_out_outputs; // decoder_output + std::string dec_out_prednet_lengths; + std::string dec_out_states_1; // h_out + std::string dec_out_states_2; // c_out + + // Joiner I/O names + std::string join_in_encoder; + std::string join_in_decoder; + std::string join_out_logits; + + /// Populate from a Config object (reads encoder/decoder/joiner/speech sections). + void PopulateFromConfig(const Config& config); +}; + +/// Holds the TDT decoder LSTM hidden states between decoding steps. +struct ParakeetDecoderState { + // h_in / c_in: [lstm_layers, 1, lstm_dim] + std::unique_ptr state_h; + std::unique_ptr state_c; + + // Decoder output from last step: [1, lstm_dim, 1] + std::unique_ptr decoder_output; + + int64_t last_token{0}; // Last emitted non-blank token (for autoregressive feedback) + + void Initialize(const ParakeetConfig& cfg, OrtAllocator& allocator); + void Reset(const ParakeetConfig& cfg, OrtAllocator& allocator); + + private: + void RunInitialDecoder(const ParakeetConfig& cfg, OrtAllocator& allocator, + OrtSession* decoder_session); + friend struct ParakeetStreamingASR; +}; + +// ─── Model ────────────────────────────────────────────────────────────────── + +struct ParakeetSpeechModel : Model { + ParakeetSpeechModel(std::unique_ptr config, OrtEnv& ort_env); + + std::unique_ptr CreateState(DeviceSpan sequence_lengths, + const GeneratorParams& params) const override; + + // Three ONNX sessions: encoder, decoder (prediction network), joiner + std::unique_ptr session_encoder_; + std::unique_ptr session_decoder_; + std::unique_ptr session_joiner_; + + std::unique_ptr encoder_session_options_; + std::unique_ptr decoder_session_options_; + std::unique_ptr joiner_session_options_; + + ParakeetConfig parakeet_config_; +}; + +} // namespace Generators diff --git a/src/ort_genai.h b/src/ort_genai.h index 949dff9d96..76130c4a1e 100644 --- a/src/ort_genai.h +++ b/src/ort_genai.h @@ -918,3 +918,35 @@ struct OgaStreamingProcessor : OgaAbstract { static void operator delete(void* p) { OgaDestroyStreamingProcessor(reinterpret_cast(p)); } }; + +struct OgaStreamingASR : OgaAbstract { + static std::unique_ptr Create(OgaModel& model) { + OgaStreamingASR* p; + OgaCheckResult(OgaCreateStreamingASR(&model, &p)); + return std::unique_ptr(p); + } + + OgaString TranscribeChunk(const float* audio_data, size_t num_samples) { + const char* text; + OgaCheckResult(OgaStreamingASRTranscribeChunk(this, audio_data, num_samples, &text)); + return text; + } + + OgaString GetTranscript() const { + const char* text; + OgaCheckResult(OgaStreamingASRGetTranscript(this, &text)); + return text; + } + + void Reset() { + OgaCheckResult(OgaStreamingASRReset(this)); + } + + OgaString Flush() { + const char* text; + OgaCheckResult(OgaStreamingASRFlush(this, &text)); + return text; + } + + static void operator delete(void* p) { OgaDestroyStreamingASR(reinterpret_cast(p)); } +}; diff --git a/src/ort_genai_c.cpp b/src/ort_genai_c.cpp index 3a0156098d..bbd6822933 100644 --- a/src/ort_genai_c.cpp +++ b/src/ort_genai_c.cpp @@ -13,6 +13,7 @@ #include "search.h" #include "smartptrs.h" #include "engine/engine.h" +#include "streaming_asr.h" #include "models/streaming_processor.h" #include "models/nemotron_speech.h" #include "models/silero_vad.h" @@ -64,6 +65,7 @@ struct OgaTokenizerStream : Generators::TokenizerStream, OgaAbstract {}; struct OgaEngine : Generators::Engine, OgaAbstract {}; struct OgaRequest : Generators::Request, OgaAbstract {}; struct OgaStreamingProcessor : Generators::StreamingProcessor, OgaAbstract {}; +struct OgaStreamingASR : Generators::StreamingASR, OgaAbstract {}; // Helper function to return a shared pointer as a raw pointer. It won't compile if the types are wrong. // Exposed types that are internally owned by shared_ptrs inherit from ExternalRefCounted. Then we @@ -1095,6 +1097,46 @@ void OGA_API_CALL OgaUnregisterExecutionProviderLibrary(const char* registration Ort::UnregisterExecutionProviderLibrary(&(Generators::GetOrtEnv()), registration_name); } +OgaResult* OGA_API_CALL OgaCreateStreamingASR(OgaModel* model, OgaStreamingASR** out) { + OGA_TRY + auto asr = Generators::CreateStreamingASR(*model); + *out = ReturnUnique(std::move(asr)); + return nullptr; + OGA_CATCH +} + +OgaResult* OGA_API_CALL OgaStreamingASRTranscribeChunk(OgaStreamingASR* asr, const float* audio_data, size_t num_samples, const char** text) { + OGA_TRY + std::string result = asr->TranscribeChunk(audio_data, num_samples); + *text = AllocOgaString(result); + return nullptr; + OGA_CATCH +} + +OgaResult* OGA_API_CALL OgaStreamingASRGetTranscript(const OgaStreamingASR* asr, const char** text) { + OGA_TRY + *text = AllocOgaString(asr->GetTranscript()); + return nullptr; + OGA_CATCH +} + +OgaResult* OGA_API_CALL OgaStreamingASRReset(OgaStreamingASR* asr) { + OGA_TRY + asr->Reset(); + return nullptr; + OGA_CATCH +} + +OgaResult* OGA_API_CALL OgaStreamingASRFlush(OgaStreamingASR* asr, const char** text) { + OGA_TRY + std::string result = asr->Flush(); + *text = AllocOgaString(result); + return nullptr; + OGA_CATCH +} + +void OGA_API_CALL OgaDestroyStreamingASR(OgaStreamingASR* p) { delete p; } + OgaResult* OGA_API_CALL OgaCreateStreamingProcessor(OgaModel* model, OgaStreamingProcessor** out) { OGA_TRY auto processor = Generators::CreateStreamingProcessor(*model); diff --git a/src/ort_genai_c.h b/src/ort_genai_c.h index b507ca5b14..55bbf18028 100644 --- a/src/ort_genai_c.h +++ b/src/ort_genai_c.h @@ -79,6 +79,7 @@ typedef struct OgaStringArray OgaStringArray; typedef struct OgaAdapters OgaAdapters; typedef struct OgaEngine OgaEngine; typedef struct OgaRequest OgaRequest; +typedef struct OgaStreamingASR OgaStreamingASR; typedef struct OgaStreamingProcessor OgaStreamingProcessor; //! @} @@ -1147,6 +1148,39 @@ OGA_EXPORT void OGA_API_CALL OgaRegisterExecutionProviderLibrary(const char* reg */ OGA_EXPORT void OGA_API_CALL OgaUnregisterExecutionProviderLibrary(const char* registration_name); +/** + * \brief Creates a StreamingASR instance for real-time streaming speech recognition. + * \param[in] model The model to use for streaming ASR (must be parakeet_tdt type). + * \param[out] out Pointer to store the created StreamingASR instance. + * \return OgaResult on error, nullptr on success. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaCreateStreamingASR(OgaModel* model, OgaStreamingASR** out); + +/** + * \brief Transcribe an audio chunk. Returns newly transcribed text. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaStreamingASRTranscribeChunk(OgaStreamingASR* asr, const float* audio_data, size_t num_samples, const char** text); + +/** + * \brief Get the full transcript accumulated so far. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaStreamingASRGetTranscript(const OgaStreamingASR* asr, const char** text); + +/** + * \brief Reset streaming state for a new utterance. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaStreamingASRReset(OgaStreamingASR* asr); + +/** + * \brief Flush remaining buffered audio. Call after the last TranscribeChunk. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaStreamingASRFlush(OgaStreamingASR* asr, const char** text); + +/** + * \brief Destroy a StreamingASR instance. + */ +OGA_EXPORT void OGA_API_CALL OgaDestroyStreamingASR(OgaStreamingASR* asr); + /** * \brief Creates a StreamingProcessor for mel spectrogram extraction from raw audio. * \param[in] model The model to create the processor for (must be nemotron_speech type). diff --git a/src/parakeet_mel.cpp b/src/parakeet_mel.cpp new file mode 100644 index 0000000000..52906c89c1 --- /dev/null +++ b/src/parakeet_mel.cpp @@ -0,0 +1,249 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Parakeet mel spectrogram — standalone C++ implementation matching +// NeMo's AudioToMelSpectrogramPreprocessor + librosa.filters.mel exactly. +// +// No dependency on onnxruntime-extensions. + +#include "parakeet_mel.h" + +#include +#include +#include +#include +#include + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +namespace parakeet_mel { + +// ─── Symmetric Hann window (periodic=False) ───────────────────────────────── +// w[n] = 0.5 * (1 - cos(2*pi*n / (N-1))) for n = 0..N-1 +// This matches torch.hann_window(N, periodic=False) + +static std::vector SymmetricHannWindow(int length) { + std::vector window(length); + if (length == 1) { + window[0] = 1.0f; + return window; + } + for (int i = 0; i < length; ++i) { + window[i] = 0.5f * (1.0f - std::cos(2.0 * M_PI * i / (length - 1))); + } + return window; +} + +// ─── Radix-2 Cooley-Tukey FFT ─────────────────────────────────────────────── +// In-place, decimation-in-time. N must be a power of 2. + +static void FFT(std::vector>& x) { + int N = static_cast(x.size()); + if (N <= 1) return; + + // Bit-reversal permutation + for (int i = 1, j = 0; i < N; ++i) { + int bit = N >> 1; + for (; j & bit; bit >>= 1) { + j ^= bit; + } + j ^= bit; + if (i < j) std::swap(x[i], x[j]); + } + + // Butterfly stages + for (int len = 2; len <= N; len <<= 1) { + float angle = -2.0f * static_cast(M_PI) / len; + std::complex wlen(std::cos(angle), std::sin(angle)); + for (int i = 0; i < N; i += len) { + std::complex w(1.0f, 0.0f); + for (int j = 0; j < len / 2; ++j) { + auto u = x[i + j]; + auto v = x[i + j + len / 2] * w; + x[i + j] = u + v; + x[i + j + len / 2] = u - v; + w *= wlen; + } + } + } +} + +// ─── Librosa-compatible mel filterbank (Slaney scale) ─────────────────────── +// Matches librosa.filters.mel(sr, n_fft, n_mels, fmin, fmax, htk=False, norm=None) +// +// Slaney mel scale: +// - Linear region (f < 1000 Hz): mel = 3 * f / 200 +// - Log region (f >= 1000 Hz): mel = 15 + 27 * log(f / 1000) / log(6.4) + +static float HzToMelSlaney(float hz) { + const float f_sp = 200.0f / 3.0f; // = 66.667 Hz per mel below 1000 Hz + float mel = hz / f_sp; + + const float min_log_hz = 1000.0f; + const float min_log_mel = min_log_hz / f_sp; // = 15.0 + const float logstep = std::log(6.4f) / 27.0f; + + if (hz >= min_log_hz) { + mel = min_log_mel + std::log(hz / min_log_hz) / logstep; + } + return mel; +} + +static float MelToHzSlaney(float mel) { + const float f_sp = 200.0f / 3.0f; + float hz = mel * f_sp; + + const float min_log_hz = 1000.0f; + const float min_log_mel = min_log_hz / f_sp; + const float logstep = std::log(6.4f) / 27.0f; + + if (mel >= min_log_mel) { + hz = min_log_hz * std::exp(logstep * (mel - min_log_mel)); + } + return hz; +} + +// Build triangular mel filterbank [num_mels, num_fft_bins] +// num_fft_bins = n_fft / 2 + 1 +// This matches librosa.filters.mel with htk=False, norm=None +static std::vector> CreateMelFilterbank(int num_mels, int fft_size, + int sample_rate, + float fmin, float fmax) { + int num_bins = fft_size / 2 + 1; // 257 + + // Create num_mels + 2 mel-spaced points between fmin and fmax + float mel_min = HzToMelSlaney(fmin); + float mel_max = HzToMelSlaney(fmax); + + int num_points = num_mels + 2; + std::vector mel_points(num_points); + for (int i = 0; i < num_points; ++i) { + mel_points[i] = mel_min + (mel_max - mel_min) * i / (num_points - 1); + } + + // Convert mel points back to Hz + std::vector hz_points(num_points); + for (int i = 0; i < num_points; ++i) { + hz_points[i] = MelToHzSlaney(mel_points[i]); + } + + // Convert Hz points to FFT bin indices (fractional) + std::vector fft_freqs(num_bins); + for (int i = 0; i < num_bins; ++i) { + fft_freqs[i] = static_cast(sample_rate) * i / fft_size; + } + + // Build triangular filters + std::vector> filters(num_mels, std::vector(num_bins, 0.0f)); + + for (int m = 0; m < num_mels; ++m) { + float left = hz_points[m]; + float center = hz_points[m + 1]; + float right = hz_points[m + 2]; + + for (int k = 0; k < num_bins; ++k) { + float freq = fft_freqs[k]; + + if (freq >= left && freq <= center && center > left) { + filters[m][k] = (freq - left) / (center - left); + } else if (freq >= center && freq <= right && right > center) { + filters[m][k] = (right - freq) / (right - center); + } + } + } + + return filters; +} + +// ─── ComputeLogMel ────────────────────────────────────────────────────────── + +std::vector ComputeLogMel(const float* audio, size_t num_samples, + const ParakeetMelConfig& cfg, int& out_num_frames) { + const int n_fft = cfg.fft_size; + const int hop = cfg.hop_length; + const int win_len = cfg.win_length; + const int num_mels = cfg.num_mels; + const int num_bins = n_fft / 2 + 1; // 257 + const float log_guard = std::pow(2.0f, -24.0f); // NeMo default: 2^-24 + + // 1. Preemphasis: x[0] unchanged, x[n] = x[n] - preemph * x[n-1] + std::vector preemph_audio(num_samples); + if (num_samples > 0) { + preemph_audio[0] = audio[0]; + for (size_t i = 1; i < num_samples; ++i) { + preemph_audio[i] = audio[i] - cfg.preemph * audio[i - 1]; + } + } + + // 2. Center-pad: zero-pad n_fft/2 on each side (matches torch.stft center=True, pad_mode="constant") + int pad = n_fft / 2; // 256 + size_t padded_len = num_samples + 2 * pad; + std::vector padded(padded_len, 0.0f); + std::memcpy(padded.data() + pad, preemph_audio.data(), num_samples * sizeof(float)); + + // 3. Create symmetric Hann window, center-padded to n_fft + // torch.stft centers the window when win_length < n_fft: + // pad_left = (n_fft - win_length) / 2 + auto hann = SymmetricHannWindow(win_len); + std::vector window(n_fft, 0.0f); + int win_offset = (n_fft - win_len) / 2; + for (int i = 0; i < win_len; ++i) { + window[win_offset + i] = hann[i]; + } + + // 4. Compute STFT frames + int num_stft_frames = static_cast((padded_len - n_fft) / hop) + 1; + + // Build mel filterbank once + auto mel_filters = CreateMelFilterbank(num_mels, n_fft, cfg.sample_rate, cfg.fmin, cfg.fmax); + + // Output: [num_mels, num_valid_frames] where num_valid_frames = num_samples // hop_length + int valid_frames = static_cast(num_samples) / hop; + if (valid_frames > num_stft_frames) valid_frames = num_stft_frames; + out_num_frames = valid_frames; + + if (valid_frames <= 0) { + out_num_frames = 0; + return {}; + } + + std::vector result(num_mels * valid_frames, 0.0f); + + // FFT buffer (reused per frame) + std::vector> fft_buf(n_fft); + std::vector power_spectrum(num_bins); + + for (int frame = 0; frame < valid_frames; ++frame) { + // Window the frame + const float* frame_start = padded.data() + frame * hop; + for (int i = 0; i < n_fft; ++i) { + fft_buf[i] = std::complex(frame_start[i] * window[i], 0.0f); + } + + // FFT + FFT(fft_buf); + + // Power spectrum: |FFT|^2 (magnitude squared) + // This matches: magnitude = sqrt(re^2 + im^2), then power = magnitude^2 = re^2 + im^2 + for (int k = 0; k < num_bins; ++k) { + float re = fft_buf[k].real(); + float im = fft_buf[k].imag(); + power_spectrum[k] = re * re + im * im; + } + + // Apply mel filterbank + log + for (int m = 0; m < num_mels; ++m) { + float mel_energy = 0.0f; + for (int k = 0; k < num_bins; ++k) { + mel_energy += mel_filters[m][k] * power_spectrum[k]; + } + result[m * valid_frames + frame] = std::log(mel_energy + log_guard); + } + } + + return result; +} + +} // namespace parakeet_mel diff --git a/src/parakeet_mel.h b/src/parakeet_mel.h new file mode 100644 index 0000000000..4619307b85 --- /dev/null +++ b/src/parakeet_mel.h @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Parakeet mel spectrogram — matches NeMo's AudioToMelSpectrogramPreprocessor +// and librosa.filters.mel exactly (Slaney scale, symmetric Hann, center=True). +// +// This is a standalone implementation with no dependency on onnxruntime-extensions. +#pragma once + +#include +#include + +namespace parakeet_mel { + +struct ParakeetMelConfig { + int num_mels = 128; + int fft_size = 512; // n_fft + int hop_length = 160; + int win_length = 400; + int sample_rate = 16000; + float preemph = 0.97f; + float fmin = 0.0f; + float fmax = 8000.0f; +}; + +/// Compute log-mel spectrogram matching NeMo/Parakeet exactly. +/// +/// Pipeline (matches parakeet_onnx_streaming_continue.py): +/// 1. Preemphasis: x[0] unchanged, x[n] = x[n] - 0.97*x[n-1] +/// 2. STFT: center=True (zero-pad n_fft/2 both sides), symmetric Hann window, +/// n_fft=512, hop=160, win=400 +/// 3. Power spectrum: |STFT|^2 +/// 4. Mel filterbank: librosa-compatible (Slaney scale, 128 bins, fmin=0, fmax=8000) +/// 5. Log: log(mel + 2^-24) (NeMo default log_zero_guard) +/// 6. Truncate to seq_len // hop_length frames +/// +/// Output layout: row-major [num_mels, num_frames]. +/// out_num_frames is set to the number of time frames produced. +std::vector ComputeLogMel(const float* audio, size_t num_samples, + const ParakeetMelConfig& cfg, int& out_num_frames); + +} // namespace parakeet_mel diff --git a/src/parakeet_streaming_asr.cpp b/src/parakeet_streaming_asr.cpp new file mode 100644 index 0000000000..06a9b9cb14 --- /dev/null +++ b/src/parakeet_streaming_asr.cpp @@ -0,0 +1,532 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// ParakeetStreamingASR — streaming ASR for Parakeet FastConformer + TDT models. +// +// Implements the same streaming algorithm as NVIDIA NeMo's +// speech_to_text_streaming_infer_rnnt.py, adapted for TDT (Token-and-Duration +// Transducer) decoding where the joiner predicts both a token and a duration. +// +// Key difference from NemoStreamingASR (RNNT): +// - Non-cache-aware encoder: re-encodes [left_context | chunk | right_context] +// each iteration, with per-window mel normalization +// - TDT: joiner output has vocab_size + num_durations logits; the duration +// controls how many encoder frames to advance (not just blank=1) + +#include +#include +#include +#include +#include + +#include "generators.h" +#include "parakeet_streaming_asr.h" +#include "parakeet_mel.h" + +namespace Generators { + +// ─── Vocabulary loading ───────────────────────────────────────────────────── + +void ParakeetStreamingASR::LoadVocab() { + if (vocab_loaded_) return; + + // Try vocab.txt first (one token per line, 0-indexed); fall back to tokens.txt + // (sherpa-onnx format: " " per line). We read directly from the + // file rather than using the tokenizer's Decode() because SPM's Decode strips + // the leading ▁ space marker for individual tokens, which causes word + // boundaries to be lost during streaming concatenation. + vocab_.assign(config_.vocab_size, std::string{}); + + auto vocab_path = model_.config_->config_path / "vocab.txt"; + std::ifstream f(vocab_path.string()); + if (f.is_open()) { + std::string line; + int idx = 0; + while (std::getline(f, line) && idx < config_.vocab_size) { + while (!line.empty() && (line.back() == '\r' || line.back() == '\n')) + line.pop_back(); + vocab_[idx] = line; + idx++; + } + } else { + auto tokens_path = model_.config_->config_path / "tokens.txt"; + std::ifstream tf(tokens_path.string()); + if (tf.is_open()) { + std::string line; + while (std::getline(tf, line)) { + while (!line.empty() && (line.back() == '\r' || line.back() == '\n')) + line.pop_back(); + // Format: " " — split on the LAST space so token may contain spaces. + auto sp = line.find_last_of(' '); + if (sp == std::string::npos) continue; + std::string tok = line.substr(0, sp); + int id = std::atoi(line.c_str() + sp + 1); + if (id >= 0 && id < config_.vocab_size) vocab_[id] = std::move(tok); + } + } else { + // Final fallback: use tokenizer Decode (may lose spaces) + auto tokenizer = model_.CreateTokenizer(); + for (int i = 0; i < config_.vocab_size; ++i) { + try { + std::vector ids = {static_cast(i)}; + vocab_[i] = tokenizer->Decode(ids); + } catch (...) { + vocab_[i] = ""; + } + } + } + } + vocab_loaded_ = true; +} + +// ─── Constructor / Reset ──────────────────────────────────────────────────── + +ParakeetStreamingASR::ParakeetStreamingASR(Model& model) + : model_{model} { + auto* parakeet_model = dynamic_cast(&model); + if (!parakeet_model) { + throw std::runtime_error("ParakeetStreamingASR requires a parakeet_tdt model type. Got: " + model.config_->model.type); + } + + encoder_session_ = parakeet_model->session_encoder_.get(); + decoder_session_ = parakeet_model->session_decoder_.get(); + joiner_session_ = parakeet_model->session_joiner_.get(); + config_ = parakeet_model->parakeet_config_; + + // Initialize decoder state (LSTM zeros) + auto& allocator = model_.allocator_cpu_; + decoder_state_.Initialize(config_, allocator); + decoder_initialized_ = false; +} + +ParakeetStreamingASR::~ParakeetStreamingASR() = default; + +void ParakeetStreamingASR::Reset() { + auto& allocator = model_.allocator_cpu_; + decoder_state_.Reset(config_, allocator); + decoder_initialized_ = false; + full_transcript_.clear(); + all_audio_.clear(); + processed_audio_samples_ = 0; + chunk_index_ = 0; +} + +// ─── Initialize decoder with blank token ──────────────────────────────────── + +void ParakeetStreamingASR::InitializeDecoderState() { + if (decoder_initialized_) return; + + auto& allocator = model_.allocator_cpu_; + auto run_options = OrtRunOptions::Create(); + + // Run decoder with blank_id to get initial decoder_output + auto targets_shape = std::array{1, 1}; + auto targets = OrtValue::CreateTensor(allocator, targets_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + *targets->GetTensorMutableData() = static_cast(config_.blank_id); + + auto tgt_len_shape = std::array{1}; + auto target_length = OrtValue::CreateTensor(allocator, tgt_len_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + *target_length->GetTensorMutableData() = 1; + + const char* dec_input_names[] = { + config_.dec_in_targets.c_str(), config_.dec_in_target_length.c_str(), + config_.dec_in_states_1.c_str(), config_.dec_in_states_2.c_str()}; + OrtValue* dec_inputs[] = { + targets.get(), target_length.get(), + decoder_state_.state_h.get(), decoder_state_.state_c.get()}; + + const char* dec_output_names[] = { + config_.dec_out_outputs.c_str(), config_.dec_out_prednet_lengths.c_str(), + config_.dec_out_states_1.c_str(), config_.dec_out_states_2.c_str()}; + + auto dec_outputs = decoder_session_->Run( + run_options.get(), + dec_input_names, dec_inputs, 4, + dec_output_names, 4); + + // decoder_output shape: [1, 640, 1] — take last time step + auto dec_out_shape = dec_outputs[0]->GetTensorTypeAndShapeInfo()->GetShape(); + int64_t dec_dim = dec_out_shape[1]; + + // Extract last time step: [1, 640, 1] + auto frame_shape = std::array{1, dec_dim, 1}; + decoder_state_.decoder_output = OrtValue::CreateTensor(allocator, frame_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + const float* src = dec_outputs[0]->GetTensorData(); + float* dst = decoder_state_.decoder_output->GetTensorMutableData(); + int64_t t_last = dec_out_shape[2] - 1; + for (int64_t d = 0; d < dec_dim; ++d) { + dst[d] = src[d * dec_out_shape[2] + t_last]; + } + + decoder_state_.state_h = std::move(dec_outputs[2]); + decoder_state_.state_c = std::move(dec_outputs[3]); + + decoder_initialized_ = true; +} + +// ─── TranscribeChunk ──────────────────────────────────────────────────────── + +std::string ParakeetStreamingASR::TranscribeChunk(const float* audio_data, size_t num_samples) { + LoadVocab(); + InitializeDecoderState(); + + // Append incoming audio to full audio buffer + all_audio_.insert(all_audio_.end(), audio_data, audio_data + num_samples); + + const size_t chunk_sz = static_cast(config_.chunk_samples); + const size_t right_ctx = static_cast(config_.right_context_samples); + std::string result; + + // Process chunks only when we have enough audio for right context lookahead. + // This matches the Python reference which has the full audio available: + // win_right = min(chunk_end + right_samples, len(audio)) + // By waiting for right_ctx more samples, we ensure the encoder window + // includes the right context audio. + while (processed_audio_samples_ + chunk_sz + right_ctx <= all_audio_.size()) { + size_t chunk_start = processed_audio_samples_; + size_t chunk_end = chunk_start + chunk_sz; + bool is_last = false; // Not last — there's at least right_ctx of future audio + + result += ProcessChunk(chunk_start, chunk_end, is_last); + processed_audio_samples_ = chunk_end; + chunk_index_++; + } + + return result; +} + +std::string ParakeetStreamingASR::Flush() { + LoadVocab(); + InitializeDecoderState(); + + std::string result; + const size_t chunk_sz = static_cast(config_.chunk_samples); + + // Process any remaining complete chunks (without requiring right context) + while (processed_audio_samples_ + chunk_sz <= all_audio_.size()) { + size_t chunk_start = processed_audio_samples_; + size_t chunk_end = chunk_start + chunk_sz; + bool is_last = (chunk_end + chunk_sz > all_audio_.size()); + + result += ProcessChunk(chunk_start, chunk_end, is_last); + processed_audio_samples_ = chunk_end; + chunk_index_++; + } + + // Process final partial chunk if any audio remains + if (processed_audio_samples_ < all_audio_.size()) { + size_t chunk_start = processed_audio_samples_; + size_t chunk_end = all_audio_.size(); + + result += ProcessChunk(chunk_start, chunk_end, /*is_last=*/true); + processed_audio_samples_ = chunk_end; + chunk_index_++; + } + + return result; +} + +// ─── ProcessChunk: build window, mel, normalize, encode, TDT decode ──────── + +std::string ParakeetStreamingASR::ProcessChunk(size_t chunk_start, size_t chunk_end, bool is_last) { + auto& allocator = model_.allocator_cpu_; + + const int hop = config_.hop_length; + const int sub = config_.subsampling_factor; + const int encoder_frame_samples = hop * sub; // 1280 + + // Streaming context sizes from config (already aligned to encoder frame boundaries) + const size_t left_samples = static_cast(config_.left_context_samples); + const size_t right_samples = static_cast(config_.right_context_samples); + // Chunk size aligned + const size_t chunk_samples_aligned = static_cast( + (static_cast(chunk_end - chunk_start) / encoder_frame_samples) * encoder_frame_samples); + + // Build encoder window: [left_context | chunk | right_context] + size_t win_left = (chunk_start > left_samples) ? (chunk_start - left_samples) : 0; + size_t win_right = std::min(chunk_end + right_samples, all_audio_.size()); + + // For last chunk, try to include as much context as possible + if (is_last) { + size_t target_buf_size = left_samples + chunk_samples_aligned + right_samples; + size_t actual_size = win_right - win_left; + if (actual_size < target_buf_size) { + win_left = (win_right > target_buf_size) ? (win_right - target_buf_size) : 0; + } + } + + const float* window_audio = all_audio_.data() + win_left; + size_t window_len = win_right - win_left; + + // ── Compute mel features on the window ── + // Use parakeet_mel::ComputeLogMel which matches NeMo's pipeline exactly: + // preemphasis, symmetric Hann, center-padded STFT, librosa mel filterbank, log(x + 2^-24) + parakeet_mel::ParakeetMelConfig mel_cfg; + mel_cfg.num_mels = config_.num_mels; + mel_cfg.fft_size = config_.fft_size; + mel_cfg.hop_length = config_.hop_length; + mel_cfg.win_length = config_.win_length; + mel_cfg.sample_rate = config_.sample_rate; + mel_cfg.preemph = config_.preemph; + + int num_mel_frames = 0; + auto raw_mel = parakeet_mel::ComputeLogMel(window_audio, window_len, mel_cfg, num_mel_frames); + // raw_mel is [num_mels, num_mel_frames] row-major + + if (num_mel_frames <= 0) return ""; + + // ── Per-window normalization (Bessel's correction, per-feature) ── + // This matches NeMo's normalize_batch with per_feature mode + const int num_mels = config_.num_mels; + std::vector mel_normalized(num_mels * num_mel_frames); + + for (int m = 0; m < num_mels; ++m) { + const float* row = raw_mel.data() + m * num_mel_frames; + + // Compute mean + double sum = 0.0; + for (int t = 0; t < num_mel_frames; ++t) sum += row[t]; + float mean = static_cast(sum / num_mel_frames); + + // Compute variance with Bessel's correction + double var_sum = 0.0; + for (int t = 0; t < num_mel_frames; ++t) { + double diff = row[t] - mean; + var_sum += diff * diff; + } + float std_val; + if (num_mel_frames > 1) { + std_val = std::sqrt(static_cast(var_sum / (num_mel_frames - 1))) + 1e-5f; + } else { + std_val = 1e-5f; + } + + // Normalize + float* out_row = mel_normalized.data() + m * num_mel_frames; + for (int t = 0; t < num_mel_frames; ++t) { + out_row[t] = (row[t] - mean) / std_val; + } + } + + // ── Run encoder ── + auto signal_shape = std::array{1, static_cast(num_mels), static_cast(num_mel_frames)}; + auto processed_signal = OrtValue::CreateTensor(allocator, signal_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + std::memcpy(processed_signal->GetTensorMutableData(), mel_normalized.data(), + num_mels * num_mel_frames * sizeof(float)); + + auto len_shape = std::array{1}; + auto signal_length = OrtValue::CreateTensor(allocator, len_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + *signal_length->GetTensorMutableData() = static_cast(num_mel_frames); + + const char* enc_input_names[] = { + config_.enc_in_audio.c_str(), config_.enc_in_length.c_str()}; + OrtValue* enc_inputs[] = {processed_signal.get(), signal_length.get()}; + + const char* enc_output_names[] = { + config_.enc_out_encoded.c_str(), config_.enc_out_length.c_str()}; + + auto run_options = OrtRunOptions::Create(); + auto enc_outputs = encoder_session_->Run( + run_options.get(), + enc_input_names, enc_inputs, 2, + enc_output_names, 2); + + auto* encoded = enc_outputs[0].get(); + int64_t enc_total = *enc_outputs[1]->GetTensorData(); + + // ── Strip left context encoder frames ── + size_t left_ctx_samples = chunk_start - win_left; + int64_t left_ctx_mel = static_cast(left_ctx_samples) / hop; + int64_t left_enc = left_ctx_mel / sub; + + // ── Determine decode range (only decode chunk frames, not left/right context) ── + int64_t decode_start = left_enc; + int64_t decode_end; + + if (is_last) { + decode_end = enc_total; + } else { + int64_t chunk_actual = static_cast(chunk_end - chunk_start); + int64_t chunk_mel = chunk_actual / hop; + int64_t chunk_enc = chunk_mel / sub; + decode_end = std::min(left_enc + chunk_enc, enc_total); + } + + if (decode_end <= decode_start) return ""; + + // ── TDT greedy decode ── + std::string chunk_text = RunTDTDecoder(encoded, enc_total, decode_start, decode_end); + full_transcript_ += chunk_text; + + return chunk_text; +} + +// ─── TDT Greedy Decoder ───────────────────────────────────────────────────── + +std::string ParakeetStreamingASR::RunTDTDecoder(OrtValue* encoder_output, + int64_t encoded_len, + int64_t start_frame, + int64_t end_frame) { + auto& allocator = model_.allocator_cpu_; + std::string result; + + auto enc_info = encoder_output->GetTensorTypeAndShapeInfo(); + auto enc_shape = enc_info->GetShape(); + // enc_shape: [1, 1024, T'] — [batch, hidden_dim, time] + int64_t hidden_dim = enc_shape[1]; + int64_t enc_time = enc_shape[2]; + const float* enc_data = encoder_output->GetTensorData(); + + auto run_options = OrtRunOptions::Create(); + + const int num_durations = config_.tdt_num_extra_outputs; + const int vocab_size = config_.vocab_size; + const int blank_id = config_.blank_id; + const int max_sym = config_.max_symbols_per_step; + + int symbols_this_frame = 0; + int64_t t = start_frame; + + while (t < end_frame) { + // ── Extract single encoder frame: [1, hidden_dim, 1] ── + auto frame_shape = std::array{1, hidden_dim, 1}; + auto encoder_frame_raw = OrtValue::CreateTensor(allocator, frame_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + float* frame_data = encoder_frame_raw->GetTensorMutableData(); + for (int64_t d = 0; d < hidden_dim; ++d) { + frame_data[d] = enc_data[d * enc_time + t]; + } + + // ── Joiner expects encoder_outputs [1, hidden_dim, 1] and decoder_outputs [1, dec_dim, 1] ── + auto dec_shape = decoder_state_.decoder_output->GetTensorTypeAndShapeInfo()->GetShape(); + int64_t dec_dim = dec_shape[1]; + + // ── Run joiner → [1, 1, 1, vocab_size + num_durations] ── + const char* join_input_names[] = { + config_.join_in_encoder.c_str(), config_.join_in_decoder.c_str()}; + OrtValue* join_inputs[] = {encoder_frame_raw.get(), decoder_state_.decoder_output.get()}; + + const char* join_output_names[] = {config_.join_out_logits.c_str()}; + + auto join_outputs = joiner_session_->Run( + run_options.get(), + join_input_names, join_inputs, 2, + join_output_names, 1); + + const float* logits_data = join_outputs[0]->GetTensorData(); + + // ── Token prediction: argmax over vocab + blank logits (vocab_size + 1) ── + const int num_tok_logits = vocab_size + 1; // includes blank at index vocab_size + int best_token = 0; + float best_score = logits_data[0]; + for (int i = 1; i < num_tok_logits; ++i) { + if (logits_data[i] > best_score) { + best_score = logits_data[i]; + best_token = i; + } + } + + // ── Duration prediction: argmax over next num_durations logits ── + int skip = 0; + if (num_durations > 0) { + const int dur_off = num_tok_logits; + float best_dur_score = logits_data[dur_off]; + for (int i = 1; i < num_durations; ++i) { + if (logits_data[dur_off + i] > best_dur_score) { + best_dur_score = logits_data[dur_off + i]; + skip = i; + } + } + // Map to actual duration value if available + if (skip < static_cast(config_.tdt_durations.size())) { + skip = config_.tdt_durations[skip]; + } + } + + if (best_token != blank_id) { + // ── Emit token ── + symbols_this_frame++; + + if (best_token < static_cast(vocab_.size())) { + std::string token_str = vocab_[best_token]; + // Replace sentencepiece space marker "▁" with space + size_t pos = 0; + while ((pos = token_str.find("\xe2\x96\x81", pos)) != std::string::npos) { + token_str.replace(pos, 3, " "); + pos += 1; + } + result += token_str; + } + + // ── Update decoder state with emitted token ── + auto targets_shape = std::array{1, 1}; + auto targets = OrtValue::CreateTensor(allocator, targets_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + *targets->GetTensorMutableData() = static_cast(best_token); + + auto tgt_len_shape = std::array{1}; + auto target_length = OrtValue::CreateTensor(allocator, tgt_len_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + *target_length->GetTensorMutableData() = 1; + + const char* dec_input_names[] = { + config_.dec_in_targets.c_str(), config_.dec_in_target_length.c_str(), + config_.dec_in_states_1.c_str(), config_.dec_in_states_2.c_str()}; + OrtValue* dec_inputs[] = { + targets.get(), target_length.get(), + decoder_state_.state_h.get(), decoder_state_.state_c.get()}; + + const char* dec_output_names[] = { + config_.dec_out_outputs.c_str(), config_.dec_out_prednet_lengths.c_str(), + config_.dec_out_states_1.c_str(), config_.dec_out_states_2.c_str()}; + + auto dec_outputs = decoder_session_->Run( + run_options.get(), + dec_input_names, dec_inputs, 4, + dec_output_names, 4); + + // Update decoder output: extract last time step [1, dec_dim, 1] + auto new_dec_shape = dec_outputs[0]->GetTensorTypeAndShapeInfo()->GetShape(); + int64_t new_dec_dim = new_dec_shape[1]; + int64_t new_dec_time = new_dec_shape[2]; + auto new_frame_shape = std::array{1, new_dec_dim, 1}; + decoder_state_.decoder_output = OrtValue::CreateTensor(allocator, new_frame_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + { + const float* src = dec_outputs[0]->GetTensorData(); + float* dst = decoder_state_.decoder_output->GetTensorMutableData(); + int64_t t_last = new_dec_time - 1; + for (int64_t d = 0; d < new_dec_dim; ++d) { + dst[d] = src[d * new_dec_time + t_last]; + } + } + + decoder_state_.state_h = std::move(dec_outputs[2]); + decoder_state_.state_c = std::move(dec_outputs[3]); + decoder_state_.last_token = static_cast(best_token); + } + + // Handle duration: skip > 0 means advance by 'skip' frames + if (skip > 0) { + symbols_this_frame = 0; + } + + // Safety: force advance if too many symbols at one frame + if (symbols_this_frame >= max_sym) { + symbols_this_frame = 0; + skip = 1; + } + + // Force advance if blank with duration 0 (prevent infinite loop) + if (best_token == blank_id && skip == 0) { + symbols_this_frame = 0; + skip = 1; + } + + t += skip; + } + + return result; +} + +std::unique_ptr CreateStreamingASR(Model& model) { + return std::make_unique(model); +} + +} // namespace Generators diff --git a/src/parakeet_streaming_asr.h b/src/parakeet_streaming_asr.h new file mode 100644 index 0000000000..2e4f89bb12 --- /dev/null +++ b/src/parakeet_streaming_asr.h @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// ParakeetStreamingASR — streaming ASR for Parakeet FastConformer + TDT models. +#pragma once + +#include "streaming_asr.h" +#include "models/parakeet_speech.h" + +namespace Generators { + +/// Streaming ASR implementation for Parakeet FastConformer encoder (non-cache-aware) +/// with TDT (Token-and-Duration Transducer) greedy decoding. +/// +/// Unlike RNNT which always advances by 1 encoder frame on blank, TDT predicts +/// a duration that controls how many encoder frames to advance after emitting a token. +/// +/// Streaming algorithm (matches NeMo's speech_to_text_streaming_infer_rnnt.py): +/// - Buffer = [left_context | chunk | right_context] +/// - Encoder runs on full buffer each iteration +/// - Left context encoder frames are STRIPPED — only chunk frames are decoded +/// - Decoder LSTM state carries forward between chunks +/// - Each token is decoded exactly once (TDT durations handle frame advancement) +struct ParakeetStreamingASR : StreamingASR { + explicit ParakeetStreamingASR(Model& model); + ~ParakeetStreamingASR() override; + + std::string TranscribeChunk(const float* audio_data, size_t num_samples) override; + std::string Flush() override; + const std::string& GetTranscript() const override { return full_transcript_; } + void Reset() override; + + private: + Model& model_; + ParakeetConfig config_; + + // ONNX sessions (borrowed from ParakeetSpeechModel) + OrtSession* encoder_session_{}; + OrtSession* decoder_session_{}; + OrtSession* joiner_session_{}; + + // Streaming decoder state (LSTM h/c + decoder_output maintained across chunks) + ParakeetDecoderState decoder_state_; + bool decoder_initialized_{false}; + std::string full_transcript_; + + // Vocabulary + std::vector vocab_; + bool vocab_loaded_{false}; + + // Audio accumulation buffer — all audio received so far + std::vector all_audio_; + + // How many audio samples we have already processed (decoded) in previous chunks + size_t processed_audio_samples_{0}; + + int chunk_index_{0}; + + void LoadVocab(); + void InitializeDecoderState(); + + // Process a single streaming chunk: build [left|chunk|right] window, + // compute mel, normalize, encode, decode chunk frames with TDT. + std::string ProcessChunk(size_t chunk_start, size_t chunk_end, bool is_last); + + // TDT greedy decode on encoder frames [start_frame, end_frame). + // Returns decoded text. Updates decoder_state_ in place. + std::string RunTDTDecoder(OrtValue* encoder_output, int64_t encoded_len, + int64_t start_frame, int64_t end_frame); +}; + +} // namespace Generators diff --git a/src/python/python.cpp b/src/python/python.cpp index 4dd1889104..7d45ac4abb 100644 --- a/src/python/python.cpp +++ b/src/python/python.cpp @@ -677,6 +677,27 @@ PYBIND11_MODULE(onnxruntime_genai, m) { pybind11::arg("key"), "Get a processor option value by key."); + pybind11::class_(m, "StreamingASR") + .def(pybind11::init([](OgaModel& model) { return OgaStreamingASR::Create(model); }), + "Create a StreamingASR instance for real-time streaming speech recognition.\n" + "The model must be of type 'parakeet_tdt'.") + .def("transcribe_chunk", [](OgaStreamingASR& asr, pybind11::array_t audio_chunk) -> std::string { + auto buf = audio_chunk.request(); + auto result = asr.TranscribeChunk(static_cast(buf.ptr), static_cast(buf.size)); + return std::string(result.p_); + }, pybind11::arg("audio_chunk"), + "Feed a chunk of float32 PCM audio (mono, 16kHz) and get newly transcribed text.") + .def("flush", [](OgaStreamingASR& asr) -> std::string { + auto result = asr.Flush(); + return std::string(result.p_); + }, "Flush remaining buffered audio. Call after the last transcribe_chunk.") + .def("get_transcript", [](const OgaStreamingASR& asr) -> std::string { + auto result = asr.GetTranscript(); + return std::string(result.p_); + }, "Get the full transcript accumulated so far.") + .def("reset", [](OgaStreamingASR& asr) { asr.Reset(); }, + "Reset streaming state for a new utterance."); + m.def("set_log_options", &SetLogOptions); m.def("set_log_callback", &SetLogCallback); diff --git a/src/streaming_asr.h b/src/streaming_asr.h new file mode 100644 index 0000000000..ea0956f508 --- /dev/null +++ b/src/streaming_asr.h @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// StreamingASR — abstract interface for streaming speech recognition. +// Concrete implementations live in model-specific headers (e.g. nemo_streaming_asr.h). +#pragma once + +#include "models/model.h" + +namespace Generators { + +/// Abstract base class for streaming ASR. +/// +/// Concrete implementations (e.g. NemoStreamingASR for FastConformer + RNNT) +/// handle model-specific encoder caching, mel extraction, and decoder logic. +/// +/// Usage: +/// auto model = CreateModel(env, "path/to/model"); +/// auto asr = CreateStreamingASR(*model); +/// std::string text = asr->TranscribeChunk(audio_data, num_samples); +/// std::string full = asr->GetTranscript(); +/// asr->Reset(); +/// +struct StreamingASR : LeakChecked { + virtual ~StreamingASR() = default; + + /// Feed a chunk of raw PCM audio (mono, float32, sample rate depends on model). + /// Returns newly transcribed text from this call. + virtual std::string TranscribeChunk(const float* audio_data, size_t num_samples) = 0; + + /// Flush remaining buffered audio (call after last TranscribeChunk). + /// Returns final transcribed text. + virtual std::string Flush() = 0; + + /// Get the full transcript accumulated so far. + virtual const std::string& GetTranscript() const = 0; + + /// Reset all streaming state for a new utterance. + virtual void Reset() = 0; +}; + +/// Factory: creates the appropriate StreamingASR implementation based on model type. +std::unique_ptr CreateStreamingASR(Model& model); + +} // namespace Generators From 6d88f75c6bfe51ccee0e899f87b19a93f046e617 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Thu, 7 May 2026 21:36:52 +0000 Subject: [PATCH 02/34] More fixes --- cmake/deps.txt | 2 +- src/config.cpp | 91 ++++++++-------- src/parakeet_mel.cpp | 243 +++++++++++-------------------------------- src/parakeet_mel.h | 9 +- 4 files changed, 120 insertions(+), 225 deletions(-) diff --git a/cmake/deps.txt b/cmake/deps.txt index 6920cabc7c..01397ef155 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -14,7 +14,7 @@ pybind11;https://github.com/pybind/pybind11/archive/refs/tags/v2.13.6.zip;f78029 googletest;https://github.com/google/googletest/archive/530d5c8c84abd2a46f38583ee817743c9b3a42b4.zip;5e3a61db2aa975cfd0f97ba92c818744e7fa7034 microsoft_wil;https://github.com/microsoft/wil/archive/refs/tags/v1.0.230629.1.zip;e4a542a323c070376f7c2d1973d0f7ddbc1d2fa5 directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v1.613.1.zip;47653509a3371eabb156360f42faf582f314bf2e -onnxruntime_extensions;https://github.com/microsoft/onnxruntime-extensions.git;e094cc816679d0b2b5fe2b4fd7f73e5b1844b425 +onnxruntime_extensions;https://github.com/microsoft/onnxruntime-extensions.git;aad7b671cbfc0b7af5c3b35283bdb6c6742b61c7 # These two dependencies are for the optional constrained decoding feature (USE_GUIDANCE) llguidance;https://github.com/microsoft/llguidance.git;94fa39128ef184ffeda33845f6d333f332a34b4d diff --git a/src/config.cpp b/src/config.cpp index 1df8b90366..6355118ab0 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -876,7 +876,7 @@ struct SpeechOutputs_Element : JSON::Element { }; struct Speech_Element : JSON::Element { - explicit Speech_Element(Config::Model::Speech& v) : v_{v} {} + Speech_Element(Config::Model::Speech& v, Config::Model& model) : v_{v}, model_{model} {} void OnValue(std::string_view name, JSON::Value value) override { if (name == "filename") { @@ -885,11 +885,55 @@ struct Speech_Element : JSON::Element { v_.config_filename = JSON::Get(value); } else if (name == "adapter_filename") { v_.adapter_filename = JSON::Get(value); + } else if (name == "num_mels") { + model_.num_mels = static_cast(JSON::Get(value)); + } else if (name == "fft_size") { + model_.fft_size = static_cast(JSON::Get(value)); + } else if (name == "hop_length") { + model_.hop_length = static_cast(JSON::Get(value)); + } else if (name == "win_length") { + model_.win_length = static_cast(JSON::Get(value)); + } else if (name == "preemph") { + model_.preemph = static_cast(JSON::Get(value)); + } else if (name == "log_eps") { + model_.log_eps = static_cast(JSON::Get(value)); + } else if (name == "subsampling_factor") { + model_.subsampling_factor = static_cast(JSON::Get(value)); + } else if (name == "left_context") { + model_.left_context = static_cast(JSON::Get(value)); + } else if (name == "conv_context") { + model_.conv_context = static_cast(JSON::Get(value)); + } else if (name == "pre_encode_cache_size") { + model_.pre_encode_cache_size = static_cast(JSON::Get(value)); + } else if (name == "sample_rate") { + model_.sample_rate = static_cast(JSON::Get(value)); + } else if (name == "chunk_samples") { + model_.chunk_samples = static_cast(JSON::Get(value)); + } else if (name == "blank_id") { + model_.blank_id = static_cast(JSON::Get(value)); + } else if (name == "max_symbols_per_step") { + model_.max_symbols_per_step = static_cast(JSON::Get(value)); + } else if (name == "left_context_samples") { + model_.left_context_samples = static_cast(JSON::Get(value)); + } else if (name == "right_context_samples") { + model_.right_context_samples = static_cast(JSON::Get(value)); + } else if (name == "tdt_num_extra_outputs") { + model_.tdt_num_extra_outputs = static_cast(JSON::Get(value)); + } else if (name == "enc_in_length") { + model_.enc_in_length = JSON::Get(value); + } else if (name == "enc_out_length") { + model_.enc_out_length = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } } + Element& OnArray(std::string_view name) override { + if (name == "tdt_durations") + return tdt_durations_; + throw JSON::unknown_value_error{}; + } + Element& OnObject(std::string_view name) override { if (name == "session_options") { v_.session_options = Config::SessionOptions{}; @@ -912,10 +956,12 @@ struct Speech_Element : JSON::Element { private: Config::Model::Speech& v_; + Config::Model& model_; std::unique_ptr session_options_; std::unique_ptr run_options_; SpeechInputs_Element inputs_{v_.inputs}; SpeechOutputs_Element outputs_{v_.outputs}; + Int_Array_Element tdt_durations_{model_.tdt_durations}; }; struct JoinerInputs_Element : JSON::Element { @@ -1129,44 +1175,6 @@ struct Model_Element : JSON::Element { v_.video_token_id = static_cast(JSON::Get(value)); } else if (name == "vision_start_token_id") { v_.vision_start_token_id = static_cast(JSON::Get(value)); - } else if (name == "num_mels") { - v_.num_mels = static_cast(JSON::Get(value)); - } else if (name == "fft_size") { - v_.fft_size = static_cast(JSON::Get(value)); - } else if (name == "hop_length") { - v_.hop_length = static_cast(JSON::Get(value)); - } else if (name == "win_length") { - v_.win_length = static_cast(JSON::Get(value)); - } else if (name == "preemph") { - v_.preemph = static_cast(JSON::Get(value)); - } else if (name == "log_eps") { - v_.log_eps = static_cast(JSON::Get(value)); - } else if (name == "subsampling_factor") { - v_.subsampling_factor = static_cast(JSON::Get(value)); - } else if (name == "left_context") { - v_.left_context = static_cast(JSON::Get(value)); - } else if (name == "conv_context") { - v_.conv_context = static_cast(JSON::Get(value)); - } else if (name == "pre_encode_cache_size") { - v_.pre_encode_cache_size = static_cast(JSON::Get(value)); - } else if (name == "sample_rate") { - v_.sample_rate = static_cast(JSON::Get(value)); - } else if (name == "chunk_samples") { - v_.chunk_samples = static_cast(JSON::Get(value)); - } else if (name == "blank_id") { - v_.blank_id = static_cast(JSON::Get(value)); - } else if (name == "max_symbols_per_step") { - v_.max_symbols_per_step = static_cast(JSON::Get(value)); - } else if (name == "left_context_samples") { - v_.left_context_samples = static_cast(JSON::Get(value)); - } else if (name == "right_context_samples") { - v_.right_context_samples = static_cast(JSON::Get(value)); - } else if (name == "tdt_num_extra_outputs") { - v_.tdt_num_extra_outputs = static_cast(JSON::Get(value)); - } else if (name == "enc_in_length") { - v_.enc_in_length = JSON::Get(value); - } else if (name == "enc_out_length") { - v_.enc_out_length = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } @@ -1175,8 +1183,6 @@ struct Model_Element : JSON::Element { Element& OnArray(std::string_view name) override { if (name == "eos_token_id") return eos_token_id_; - if (name == "tdt_durations") - return tdt_durations_; throw JSON::unknown_value_error{}; } @@ -1210,10 +1216,9 @@ struct Model_Element : JSON::Element { Encoder_Element encoder_{v_.encoder}; Decoder_Element decoder_{v_.decoder}; Int_Array_Element eos_token_id_{v_.eos_token_id}; - Int_Array_Element tdt_durations_{v_.tdt_durations}; Vision_Element vision_{v_.vision}; Embedding_Element embedding_{v_.embedding}; - Speech_Element speech_{v_.speech}; + Speech_Element speech_{v_.speech, v_}; Joiner_Element joiner_{v_.joiner}; VAD_Element vad_{v_.vad}; }; diff --git a/src/parakeet_mel.cpp b/src/parakeet_mel.cpp index 52906c89c1..ac667ae5c8 100644 --- a/src/parakeet_mel.cpp +++ b/src/parakeet_mel.cpp @@ -1,29 +1,41 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. // -// Parakeet mel spectrogram — standalone C++ implementation matching -// NeMo's AudioToMelSpectrogramPreprocessor + librosa.filters.mel exactly. +// Parakeet mel spectrogram — pipeline matches NeMo's +// AudioToMelSpectrogramPreprocessor + librosa.filters.mel exactly. // -// No dependency on onnxruntime-extensions. +// The heavy-lifting pieces (Slaney mel filterbank, pre-emphasis, real FFT / +// power spectrum) are delegated to onnxruntime-extensions +// (shared/api/nemo_mel_spectrogram.h, namespace nemo_mel). Only the bits that +// are parakeet/NeMo-specific are kept here: +// +// * Symmetric Hann window (torch.hann_window(N, periodic=False)). +// The batch helper in extensions (`NemoComputeLogMelBatch`) uses a +// *periodic* Hann window, so we can't reuse it directly without changing +// numerical output. +// * Center-padded STFT framing loop with the window centered inside an +// fft_size buffer (win_offset = (n_fft - win_length) / 2). +// * Truncation to `num_samples / hop_length` valid frames. +// * log_zero_guard = 2^-24 (NeMo default). #include "parakeet_mel.h" #include #include #include -#include #include +#include "nemo_mel_spectrogram.h" // onnxruntime-extensions: shared/api + #ifndef M_PI #define M_PI 3.14159265358979323846 #endif namespace parakeet_mel { -// ─── Symmetric Hann window (periodic=False) ───────────────────────────────── -// w[n] = 0.5 * (1 - cos(2*pi*n / (N-1))) for n = 0..N-1 -// This matches torch.hann_window(N, periodic=False) - +// Symmetric Hann window: matches torch.hann_window(N, periodic=False). +// Kept locally because nemo_mel::NemoComputeLogMelBatch uses the *periodic* +// variant (sin(pi*n/N)^2), which produces different numerical output. static std::vector SymmetricHannWindow(int length) { std::vector window(length); if (length == 1) { @@ -36,170 +48,45 @@ static std::vector SymmetricHannWindow(int length) { return window; } -// ─── Radix-2 Cooley-Tukey FFT ─────────────────────────────────────────────── -// In-place, decimation-in-time. N must be a power of 2. - -static void FFT(std::vector>& x) { - int N = static_cast(x.size()); - if (N <= 1) return; - - // Bit-reversal permutation - for (int i = 1, j = 0; i < N; ++i) { - int bit = N >> 1; - for (; j & bit; bit >>= 1) { - j ^= bit; - } - j ^= bit; - if (i < j) std::swap(x[i], x[j]); - } - - // Butterfly stages - for (int len = 2; len <= N; len <<= 1) { - float angle = -2.0f * static_cast(M_PI) / len; - std::complex wlen(std::cos(angle), std::sin(angle)); - for (int i = 0; i < N; i += len) { - std::complex w(1.0f, 0.0f); - for (int j = 0; j < len / 2; ++j) { - auto u = x[i + j]; - auto v = x[i + j + len / 2] * w; - x[i + j] = u + v; - x[i + j + len / 2] = u - v; - w *= wlen; - } - } - } -} - -// ─── Librosa-compatible mel filterbank (Slaney scale) ─────────────────────── -// Matches librosa.filters.mel(sr, n_fft, n_mels, fmin, fmax, htk=False, norm=None) -// -// Slaney mel scale: -// - Linear region (f < 1000 Hz): mel = 3 * f / 200 -// - Log region (f >= 1000 Hz): mel = 15 + 27 * log(f / 1000) / log(6.4) - -static float HzToMelSlaney(float hz) { - const float f_sp = 200.0f / 3.0f; // = 66.667 Hz per mel below 1000 Hz - float mel = hz / f_sp; - - const float min_log_hz = 1000.0f; - const float min_log_mel = min_log_hz / f_sp; // = 15.0 - const float logstep = std::log(6.4f) / 27.0f; - - if (hz >= min_log_hz) { - mel = min_log_mel + std::log(hz / min_log_hz) / logstep; - } - return mel; -} - -static float MelToHzSlaney(float mel) { - const float f_sp = 200.0f / 3.0f; - float hz = mel * f_sp; - - const float min_log_hz = 1000.0f; - const float min_log_mel = min_log_hz / f_sp; - const float logstep = std::log(6.4f) / 27.0f; - - if (mel >= min_log_mel) { - hz = min_log_hz * std::exp(logstep * (mel - min_log_mel)); - } - return hz; -} - -// Build triangular mel filterbank [num_mels, num_fft_bins] -// num_fft_bins = n_fft / 2 + 1 -// This matches librosa.filters.mel with htk=False, norm=None -static std::vector> CreateMelFilterbank(int num_mels, int fft_size, - int sample_rate, - float fmin, float fmax) { - int num_bins = fft_size / 2 + 1; // 257 - - // Create num_mels + 2 mel-spaced points between fmin and fmax - float mel_min = HzToMelSlaney(fmin); - float mel_max = HzToMelSlaney(fmax); - - int num_points = num_mels + 2; - std::vector mel_points(num_points); - for (int i = 0; i < num_points; ++i) { - mel_points[i] = mel_min + (mel_max - mel_min) * i / (num_points - 1); - } - - // Convert mel points back to Hz - std::vector hz_points(num_points); - for (int i = 0; i < num_points; ++i) { - hz_points[i] = MelToHzSlaney(mel_points[i]); - } - - // Convert Hz points to FFT bin indices (fractional) - std::vector fft_freqs(num_bins); - for (int i = 0; i < num_bins; ++i) { - fft_freqs[i] = static_cast(sample_rate) * i / fft_size; - } - - // Build triangular filters - std::vector> filters(num_mels, std::vector(num_bins, 0.0f)); - - for (int m = 0; m < num_mels; ++m) { - float left = hz_points[m]; - float center = hz_points[m + 1]; - float right = hz_points[m + 2]; - - for (int k = 0; k < num_bins; ++k) { - float freq = fft_freqs[k]; - - if (freq >= left && freq <= center && center > left) { - filters[m][k] = (freq - left) / (center - left); - } else if (freq >= center && freq <= right && right > center) { - filters[m][k] = (right - freq) / (right - center); - } - } - } - - return filters; -} - -// ─── ComputeLogMel ────────────────────────────────────────────────────────── - std::vector ComputeLogMel(const float* audio, size_t num_samples, const ParakeetMelConfig& cfg, int& out_num_frames) { const int n_fft = cfg.fft_size; const int hop = cfg.hop_length; const int win_len = cfg.win_length; const int num_mels = cfg.num_mels; - const int num_bins = n_fft / 2 + 1; // 257 + const int num_bins = n_fft / 2 + 1; const float log_guard = std::pow(2.0f, -24.0f); // NeMo default: 2^-24 - // 1. Preemphasis: x[0] unchanged, x[n] = x[n] - preemph * x[n-1] + // 1. Pre-emphasis: y[n] = x[n] - preemph * x[n-1] + // Delegated to onnxruntime-extensions (nemo_mel::ApplyPreemphasis). std::vector preemph_audio(num_samples); if (num_samples > 0) { - preemph_audio[0] = audio[0]; - for (size_t i = 1; i < num_samples; ++i) { - preemph_audio[i] = audio[i] - cfg.preemph * audio[i - 1]; - } + nemo_mel::ApplyPreemphasis(audio, num_samples, cfg.preemph, + /*prev_sample=*/0.0f, preemph_audio.data()); } - // 2. Center-pad: zero-pad n_fft/2 on each side (matches torch.stft center=True, pad_mode="constant") - int pad = n_fft / 2; // 256 - size_t padded_len = num_samples + 2 * pad; + // 2. Center-pad: zero-pad n_fft/2 on each side + // (matches torch.stft center=True, pad_mode="constant"). + const int pad = n_fft / 2; + const size_t padded_len = num_samples + 2 * static_cast(pad); std::vector padded(padded_len, 0.0f); - std::memcpy(padded.data() + pad, preemph_audio.data(), num_samples * sizeof(float)); - - // 3. Create symmetric Hann window, center-padded to n_fft - // torch.stft centers the window when win_length < n_fft: - // pad_left = (n_fft - win_length) / 2 - auto hann = SymmetricHannWindow(win_len); - std::vector window(n_fft, 0.0f); - int win_offset = (n_fft - win_len) / 2; - for (int i = 0; i < win_len; ++i) { - window[win_offset + i] = hann[i]; + if (num_samples > 0) { + std::memcpy(padded.data() + pad, preemph_audio.data(), num_samples * sizeof(float)); } - // 4. Compute STFT frames - int num_stft_frames = static_cast((padded_len - n_fft) / hop) + 1; - - // Build mel filterbank once - auto mel_filters = CreateMelFilterbank(num_mels, n_fft, cfg.sample_rate, cfg.fmin, cfg.fmax); - - // Output: [num_mels, num_valid_frames] where num_valid_frames = num_samples // hop_length + // 3. Symmetric Hann window, centered inside an fft_size buffer. + // torch.stft centers the window when win_length < n_fft: + // win_offset = (n_fft - win_length) / 2 + // We advance the frame pointer by win_offset and pass the unpadded + // window of length win_len (matches the pattern used in extensions' + // own NemoComputeLogMelBatch). + auto window = SymmetricHannWindow(win_len); + const int win_offset = (n_fft - win_len) / 2; + + // 4. Frame layout & truncation. + const int num_stft_frames = padded_len >= static_cast(n_fft) + ? static_cast((padded_len - n_fft) / hop) + 1 + : 0; int valid_frames = static_cast(num_samples) / hop; if (valid_frames > num_stft_frames) valid_frames = num_stft_frames; out_num_frames = valid_frames; @@ -209,35 +96,31 @@ std::vector ComputeLogMel(const float* audio, size_t num_samples, return {}; } - std::vector result(num_mels * valid_frames, 0.0f); - - // FFT buffer (reused per frame) - std::vector> fft_buf(n_fft); - std::vector power_spectrum(num_bins); + // 5. Mel filterbank (Slaney scale, librosa-compatible). + // Delegated to onnxruntime-extensions (nemo_mel::CreateMelFilterbank). + // Note: extensions builds the filterbank with fmin=0, fmax=sample_rate/2. + // Parakeet's default config has fmin=0 and fmax=sample_rate/2 (e.g. + // fmax=8000 at sr=16000), so this is identical. + auto mel_filters = + nemo_mel::CreateMelFilterbank(num_mels, n_fft, cfg.sample_rate); + + // 6. Per-frame STFT power spectrum + mel projection + log. + // The real-FFT power spectrum is computed by extensions + // (nemo_mel::ComputeSTFTFrame, backed by dlib::fftr). + std::vector result(static_cast(num_mels) * valid_frames, 0.0f); + std::vector power_spectrum; + power_spectrum.reserve(num_bins); for (int frame = 0; frame < valid_frames; ++frame) { - // Window the frame - const float* frame_start = padded.data() + frame * hop; - for (int i = 0; i < n_fft; ++i) { - fft_buf[i] = std::complex(frame_start[i] * window[i], 0.0f); - } - - // FFT - FFT(fft_buf); - - // Power spectrum: |FFT|^2 (magnitude squared) - // This matches: magnitude = sqrt(re^2 + im^2), then power = magnitude^2 = re^2 + im^2 - for (int k = 0; k < num_bins; ++k) { - float re = fft_buf[k].real(); - float im = fft_buf[k].imag(); - power_spectrum[k] = re * re + im * im; - } + const float* frame_start = padded.data() + frame * hop + win_offset; + nemo_mel::ComputeSTFTFrame(frame_start, window.data(), win_len, n_fft, + power_spectrum); - // Apply mel filterbank + log for (int m = 0; m < num_mels; ++m) { + const auto& filter = mel_filters[m]; float mel_energy = 0.0f; for (int k = 0; k < num_bins; ++k) { - mel_energy += mel_filters[m][k] * power_spectrum[k]; + mel_energy += filter[k] * power_spectrum[k]; } result[m * valid_frames + frame] = std::log(mel_energy + log_guard); } diff --git a/src/parakeet_mel.h b/src/parakeet_mel.h index 4619307b85..50720c4ae0 100644 --- a/src/parakeet_mel.h +++ b/src/parakeet_mel.h @@ -4,7 +4,14 @@ // Parakeet mel spectrogram — matches NeMo's AudioToMelSpectrogramPreprocessor // and librosa.filters.mel exactly (Slaney scale, symmetric Hann, center=True). // -// This is a standalone implementation with no dependency on onnxruntime-extensions. +// Implementation delegates the heavy lifting (Slaney mel filterbank, +// pre-emphasis, real FFT / power spectrum) to onnxruntime-extensions +// (nemo_mel::CreateMelFilterbank, nemo_mel::ApplyPreemphasis, +// nemo_mel::ComputeSTFTFrame). The pipeline orchestration (symmetric Hann, +// center-padded framing, frame-count truncation, log_zero_guard = 2^-24) is +// kept here because the corresponding `nemo_mel::NemoComputeLogMelBatch` +// helper uses a *periodic* Hann window, which would change the numerical +// output relative to the NeMo/Parakeet reference. #pragma once #include From 8ae481c075608b88e39946b9aa6edea2b30e8867 Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Thu, 7 May 2026 22:20:40 +0000 Subject: [PATCH 03/34] remove streaming asr for parakeet --- examples/python/parakeet.py | 82 ++++ examples/python/parakeet_simulate_mic.py | 116 ----- src/generators.cpp | 1 - src/leakcheck.h | 3 +- src/models/model.cpp | 6 +- src/models/model.h | 1 + src/models/parakeet.cpp | 476 ++++++++++++++++++++ src/models/parakeet.h | 159 +++++++ src/models/parakeet_processor.cpp | 86 ++++ src/models/parakeet_processor.h | 35 ++ src/models/parakeet_speech.cpp | 150 ------- src/models/parakeet_speech.h | 117 ----- src/ort_genai.h | 32 -- src/ort_genai_c.cpp | 42 -- src/ort_genai_c.h | 34 -- src/parakeet_streaming_asr.cpp | 532 ----------------------- src/parakeet_streaming_asr.h | 72 --- src/python/python.cpp | 21 - src/streaming_asr.h | 45 -- 19 files changed, 844 insertions(+), 1166 deletions(-) create mode 100644 examples/python/parakeet.py delete mode 100644 examples/python/parakeet_simulate_mic.py create mode 100644 src/models/parakeet.cpp create mode 100644 src/models/parakeet.h create mode 100644 src/models/parakeet_processor.cpp create mode 100644 src/models/parakeet_processor.h delete mode 100644 src/models/parakeet_speech.cpp delete mode 100644 src/models/parakeet_speech.h delete mode 100644 src/parakeet_streaming_asr.cpp delete mode 100644 src/parakeet_streaming_asr.h delete mode 100644 src/streaming_asr.h diff --git a/examples/python/parakeet.py b/examples/python/parakeet.py new file mode 100644 index 0000000000..a32d42e740 --- /dev/null +++ b/examples/python/parakeet.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +""" +Parakeet TDT speech recognition — Whisper-style API. + +Mirrors examples/python/whisper.py so the same public API +(`og.Audios`, `model.create_multimodal_processor()`, `og.Generator`) is used: + + python parakeet.py --model_path --audio_file