Nvidia Parakeet Tdt ASR support

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;b62dd46f0d58a2f307d17f0a430cb14051349ac6
+onnxruntime_extensions;https://github.com/microsoft/onnxruntime-extensions.git;e1cccaa7d5bd04d5d25f8c61c40cc1f4ef5ead7b
 
 # These two dependencies are for the optional constrained decoding feature (USE_GUIDANCE)
 llguidance;https://github.com/microsoft/llguidance.git;94fa39128ef184ffeda33845f6d333f332a34b4d

examples/python/parakeet.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+"""
+Parakeet TDT speech recognition
+
+  python parakeet.py --model_path <model_dir> --audio_file <audio>
+
+The model loads the audio in one shot and decodes it with the standard
+Generator loop.
+"""
+
+import argparse
+import os
+import time
+import wave
+
+import onnxruntime_genai as og
+
+
+def _audio_duration_seconds(path: str) -> float:
+    try:
+        with wave.open(path, "rb") as wf:
+            frames = wf.getnframes()
+            rate = wf.getframerate()
+            if rate > 0:
+                return frames / float(rate)
+    except wave.Error:
+        # Not a valid/parsable WAV stream; fall back to soundfile below.
+        pass
+    try:
+        import soundfile as sf  # type: ignore
+        info = sf.info(path)
+        return float(info.frames) / float(info.samplerate)
+    except Exception:
+        # soundfile missing or file unreadable; duration is non-essential
+        # for this sample, so return 0.0 rather than failing the run.
+        return 0.0
+
+
+def run(args: argparse.Namespace) -> None:
+    print("Loading model...")
+    config = og.Config(args.model_path)
+    if args.execution_provider != "follow_config":
+        config.clear_providers()
+        if args.execution_provider != "cpu":
+            print(f"Setting model to {args.execution_provider}")
+            config.append_provider(args.execution_provider)
+    model = og.Model(config)
+    processor = model.create_multimodal_processor()
+
+    if not os.path.exists(args.audio_file):
+        raise FileNotFoundError(f"Audio file not found: {args.audio_file}")
+
+    print(f"Loading audio: {args.audio_file}")
+    audios = og.Audios.open(args.audio_file)
+    audio_seconds = _audio_duration_seconds(args.audio_file)
+
+    print("Processing audio...")
+    t0 = time.perf_counter()
+    inputs = processor("", audios=audios)
+
+    params = og.GeneratorParams(model)
+
+    generator = og.Generator(model, params)
+    generator.set_inputs(inputs)
+
+    while not generator.is_done():
+        generator.generate_next_token()
+    elapsed = time.perf_counter() - t0
+
+    transcription = processor.decode(generator.get_sequence(0))
+
+    print()
+    print("Transcription:")
+    print(f"    {transcription.strip()}")
+
+    print()
+    if audio_seconds > 0:
+        rtfx = audio_seconds / elapsed if elapsed > 0 else float("inf")
+        print(f"Audio duration: {audio_seconds:.2f}s | Inference: {elapsed:.2f}s | RTFx: {rtfx:.2f}x")
+    else:
+        print(f"Inference: {elapsed:.2f}s")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-m", "--model_path", type=str, required=True, help="Path to the Parakeet model directory")
+    parser.add_argument("-a", "--audio_file", type=str, required=True, help="Path to the audio file (WAV/MP3/...)")
+    parser.add_argument(
+        "-e",
+        "--execution_provider",
+        type=str,
+        required=False,
+        default="follow_config",
+        choices=["cpu", "cuda", "follow_config"],
+        help="Execution provider. Defaults to follow_config.",
+    )
+    args = parser.parse_args()
+    run(args)

src/config.cpp

@@ -340,6 +340,8 @@ struct DecoderInputs_Element : JSON::Element {
       v_.lstm_hidden_state = JSON::Get<std::string_view>(value);
     } else if (name == "lstm_cell_state") {
       v_.lstm_cell_state = JSON::Get<std::string_view>(value);
+    } else if (name == "targets_length") {
+      v_.targets_length = JSON::Get<std::string_view>(value);
     } else {
       throw JSON::unknown_value_error{};
     }
@@ -373,6 +375,8 @@ struct DecoderOutputs_Element : JSON::Element {
       v_.lstm_hidden_state = JSON::Get<std::string_view>(value);
     } else if (name == "lstm_cell_state") {
       v_.lstm_cell_state = JSON::Get<std::string_view>(value);
+    } else if (name == "outputs_length") {
+      v_.outputs_length = JSON::Get<std::string_view>(value);
     } else {
       throw JSON::unknown_value_error{};
     }
@@ -1129,6 +1133,8 @@ struct Model_Element : JSON::Element {
       v_.preemph = static_cast<float>(JSON::Get<double>(value));
     } else if (name == "log_eps") {
       v_.log_eps = static_cast<float>(JSON::Get<double>(value));
+    } else if (name == "norm_eps") {
+      v_.norm_eps = static_cast<float>(JSON::Get<double>(value));
     } else if (name == "subsampling_factor") {
       v_.subsampling_factor = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "left_context") {
@@ -1145,6 +1151,10 @@ struct Model_Element : JSON::Element {
       v_.blank_id = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "max_symbols_per_step") {
       v_.max_symbols_per_step = static_cast<int>(JSON::Get<double>(value));
+    } else if (name == "left_context_samples") {
+      v_.left_context_samples = static_cast<int>(JSON::Get<double>(value));
+    } else if (name == "right_context_samples") {
+      v_.right_context_samples = static_cast<int>(JSON::Get<double>(value));
     } else {
       throw JSON::unknown_value_error{};
     }
@@ -1153,6 +1163,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 +1198,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};

src/config.h

@@ -145,6 +145,7 @@ struct Config {
     int win_length{};
     float preemph{};
     float log_eps{};
+    float norm_eps{};
     int subsampling_factor{};
     int left_context{};
     int conv_context{};
@@ -154,6 +155,11 @@ struct Config {
     int blank_id{};
     int max_symbols_per_step{};
 
+    // Parakeet TDT (Token-and-Duration Transducer) parameters
+    int left_context_samples{};
+    int right_context_samples{};
+    std::vector<int> tdt_durations;  // e.g., {0, 1, 2, 3, 4}
+
     struct Encoder {
       std::string filename;
       std::optional<SessionOptions> session_options;
@@ -346,6 +352,9 @@ struct Config {
         std::string targets;
         std::string lstm_hidden_state;
         std::string lstm_cell_state;
+
+        // Parakeet TDT decoder (prediction network) extra inputs
+        std::string targets_length;
       } inputs;
 
       struct Outputs {
@@ -361,6 +370,9 @@ struct Config {
         std::string outputs;
         std::string lstm_hidden_state;
         std::string lstm_cell_state;
+
+        // Parakeet TDT decoder (prediction network) extra outputs
+        std::string outputs_length;
       } outputs;
 
       struct PipelineModel {

src/generators.cpp

@@ -6,6 +6,7 @@
 #include "generators.h"
 #include "models/streaming_processor.h"
 #include "models/nemotron_speech.h"
+#include "models/parakeet.h"
 #include "sequences.h"
 #include "models/env_utils.h"
 #include "models/model.h"
@@ -365,11 +366,11 @@ std::unique_ptr<Search> CreateSearch(const GeneratorParams& params) {
 }
 
 Generator::Generator(const Model& model, const GeneratorParams& params) : model_{model.shared_from_this()} {
-  // RNNT models don't use the traditional search/logits pipeline,
+  // RNNT and TDT models don't use the traditional search/logits pipeline,
   // so skip the standard validations and just create the state.
-  if (ModelType::IsRNNT(model.config_->model.type)) {
+  if (ModelType::IsTransducer(model.config_->model.type)) {
     state_ = model.CreateState({}, params);
-    is_nemotron_speech_model_ = dynamic_cast<NemotronSpeechState*>(state_.get()) != nullptr;
+    transducer_state_ = dynamic_cast<TransducerState*>(state_.get());
     return;
   }
 
@@ -569,18 +570,18 @@ void Generator::SetRuntimeOption(const char* key, const char* value) {
 }
 
 size_t Generator::TokenCount() const {
-  if (is_nemotron_speech_model_)
-    return static_cast<NemotronSpeechState*>(state_.get())->TokenCount();
+  if (transducer_state_)
+    return transducer_state_->TokenCount();
   return static_cast<size_t>(search_->GetSequenceLength());
 }
 
 bool Generator::IsDone() {
   ThrowErrorIfSessionTerminated(state_->session_terminated_);
 
-  if (is_nemotron_speech_model_) {
+  if (transducer_state_) {
     // Pending mel input means we haven't started processing this chunk yet
     if (!extra_inputs_.empty()) return false;
-    return static_cast<NemotronSpeechState*>(state_.get())->IsChunkDone();
+    return transducer_state_->IsChunkDone();
   }
 
   if (computed_logits_) {
@@ -613,11 +614,12 @@ void Generator::GenerateNextToken() {
 
   ThrowErrorIfSessionTerminated(state_->session_terminated_);
 
-  // RNNT models: yield one token per call from the decoder state machine
-  if (is_nemotron_speech_model_) {
+  // Transducer models (RNNT, TDT): yield one token per call by stepping
+  // the encoder/decoder/joiner loop directly.
+  if (transducer_state_) {
     state_->SetExtraInputs(extra_inputs_);
     extra_inputs_.clear();
-    static_cast<NemotronSpeechState*>(state_.get())->StepToken();
+    transducer_state_->StepToken();
     return;
   }
 

src/generators.h

@@ -40,6 +40,7 @@ void ThrowErrorIfSessionTerminated(bool is_session_terminated);
 namespace Generators {
 struct Model;
 struct State;
+struct TransducerState;
 struct Search;
 struct Tokenizer;
 struct ConstrainedLogitsProcessor;
@@ -132,7 +133,8 @@ struct Generator : LeakChecked<Generator> {
   Action last_action_{standard};
 
   // Pre-computed per-token decisions: avoid repeated checks each token
-  bool is_nemotron_speech_model_{};
+  // Non-null when the model is a transducer (RNNT, TDT); points into state_.
+  TransducerState* transducer_state_{nullptr};
   int phi3_rope_threshold_{};  // 0 means no ROPE rewind needed
   enum class SamplingMethod { kGreedy,
                               kTopK,

src/models/model.cpp

@@ -16,6 +16,8 @@
 #include "gpt.h"
 #include "decoder_only.h"
 #include "whisper.h"
+#include "parakeet.h"
+#include "parakeet_processor.h"
 #include "nemotron_speech.h"
 #include "multi_modal.h"
 #include "lfm2.h"
@@ -830,6 +832,8 @@ std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, std::unique_ptr<Config> conf
     return std::make_shared<DecoderOnly_Model>(std::move(config), ort_env);
   if (ModelType::IsRNNT(config->model.type))
     return std::make_shared<NemotronSpeechModel>(std::move(config), ort_env);
+  if (ModelType::IsTDT(config->model.type))
+    return std::make_shared<ParakeetTdtModel>(std::move(config), ort_env);
   if (ModelType::IsALM(config->model.type))
     return std::make_shared<WhisperModel>(std::move(config), ort_env);
   if (ModelType::IsVLM(config->model.type))
@@ -936,6 +940,7 @@ MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& sess
       processor_factory_{
           {"phi3v", Processor::Create<PhiImageProcessor>},
           {"whisper", Processor::Create<WhisperProcessor>},
+          {"parakeet_tdt", Processor::Create<ParakeetTdtProcessor>},
           {"phi4mm", Processor::Create<PhiMultiModalProcessor>},
           {"gemma3", Processor::Create<GemmaImageProcessor>},
           {"gemma4", Processor::Create<Gemma4MultiModalProcessor>},

src/models/model.h

@@ -9,6 +9,7 @@
 #include "utils.h"
 #include "phi_image_processor.h"
 #include "whisper_processor.h"
+#include "parakeet_processor.h"
 #include "phi_multimodal_processor.h"
 #include "gemma_image_processor.h"
 #include "gemma4_multimodal_processor.h"

src/models/model_type.h

@@ -47,6 +47,17 @@ struct ModelType {
     return std::find(rnnt_types.begin(), rnnt_types.end(), model_type) != rnnt_types.end();
   }
 
+  inline static bool IsTDT(const std::string& model_type) {
+    static constexpr std::array<std::string_view, 1> TDT = {"parakeet_tdt"};
+    return std::find(TDT.begin(), TDT.end(), model_type) != TDT.end();
+  }
+
+  // Transducer models (RNNT, TDT) bypass the standard search/logits pipeline
+  // and drive a custom encoder/decoder/joiner loop via TransducerState.
+  inline static bool IsTransducer(const std::string& model_type) {
+    return IsRNNT(model_type) || IsTDT(model_type);
+  }
+
   inline static bool IsMMM(const std::string& model_type) {
     // Multi-modal model (MMM)
     static constexpr std::array<std::string_view, 2> MMM = {"gemma4", "phi4mm"};

src/models/nemotron_speech.cpp

@@ -338,9 +338,11 @@ DeviceSpan<float> NemotronJoinerSubState::Run(int /*total_length*/, DeviceSpan<i
 
 NemotronSpeechState::NemotronSpeechState(const NemotronSpeechModel& model,
                                          const GeneratorParams& params)
-    : State{params, model},
+    : TransducerState{params, model},
       nemotron_model_{model} {
   nemotron_config_ = model.nemotron_config_;
+  // Until audio is fed via SetExtraInputs/SetInputs, the stream is idle.
+  chunk_done_ = true;
 
   encoder_state_ = std::make_unique<NemotronEncoderSubState>(model, params);
   prediction_state_ = std::make_unique<NemotronPredictionSubState>(model, params);
@@ -441,7 +443,7 @@ void NemotronSpeechState::RunEncoder() {
   current_mel_.reset();
 }
 
-std::span<const int32_t> NemotronSpeechState::StepToken() {
+void NemotronSpeechState::StepToken() {
   if (need_encoder_run_) {
     RunEncoder();
     need_encoder_run_ = false;
@@ -520,13 +522,12 @@ std::span<const int32_t> NemotronSpeechState::StepToken() {
     }
 
     last_tokens_.push_back(static_cast<int32_t>(best_token));
-    token_count_++;
-    return last_tokens_;
+    all_tokens_.push_back(static_cast<int32_t>(best_token));
+    return;
   }
 
   // Exhausted all time steps
   chunk_done_ = true;
-  return last_tokens_;
 }
 
 }  // namespace Generators