feat(tdt): add opt-in N-best beam decoding

.github/workflows/ci.yml

@@ -139,6 +139,21 @@ jobs:
             --model nvidia/parakeet-tdt_ctc-110m \
             --output /tmp/pk110m_cl.gguf
 
+      - name: Generate NeMo TDT N-best reference
+        run: |
+          .venv/bin/python scripts/gen_nemo_baseline.py \
+            --model nvidia/parakeet-tdt_ctc-110m \
+            --audio tests/fixtures/speech.wav \
+            --tdt-beam-size 4 \
+            --output /tmp/tdt_nbest_baseline.gguf
+
+      - name: Assert TDT N-best parity with NeMo
+        env:
+          PARAKEET_TEST_GGUF: /tmp/pk110m_cl.gguf
+          PARAKEET_TEST_TDT_NBEST_BASELINE: /tmp/tdt_nbest_baseline.gguf
+        run: |
+          ctest --test-dir build --output-on-failure -R '^test_tdt_beam$'
+
       - name: Transcribe speech fixture with TDT decoder
         run: |
           ./build/examples/cli/parakeet-cli transcribe \

AGENTS.md

@@ -311,6 +311,13 @@ Used by Phase 1 parity tests.  Requires the venv and a 16 kHz mono WAV.
     --model nvidia/parakeet-tdt_ctc-110m \
     --audio tests/fixtures/speech.wav \
     --output /tmp/baseline_speech.gguf
+
+# Optional TDT beam/N-best parity data:
+.venv/bin/python scripts/gen_nemo_baseline.py \
+    --model nvidia/parakeet-tdt_ctc-110m \
+    --audio tests/fixtures/speech.wav \
+    --tdt-beam-size 4 \
+    --output /tmp/tdt_nbest_baseline.gguf
 ```
 
 ## Publishing models to HuggingFace

README.md

@@ -234,6 +234,10 @@ parakeet-cli transcribe --model m.gguf --input audio.wav --timestamps
 #    "tokens":[{"id":..,"t":..,"conf":..}]}
 parakeet-cli transcribe --model m.gguf --input audio.wav --json
 
+# Offline TDT N-best hypotheses as ranked JSON
+parakeet-cli transcribe --model m.gguf --input audio.wav --decoder tdt \
+  --beam-size 4 --nbest 4
+
 # Read WAV bytes from stdin (useful with ffmpeg/curl pipelines)
 ffmpeg -i input.mp3 -f wav - | parakeet-cli transcribe --model m.gguf --input -
 
@@ -249,6 +253,10 @@ parakeet-cli transcribe --model eou.gguf --input audio.wav --stream
 
 Timestamps and confidence match NeMo's `transcribe(timestamps=True)` with the `max_prob` confidence method exactly (word offsets to 0.0 s, per-token and per-word confidence within `5e-6`), for both the TDT and CTC heads. See `docs/parity.md`. Word start and end are in seconds (`frame x hop x subsampling / sample_rate`, which works out to 0.08 s/frame here); confidence is the rescaled softmax probability of the emitted token, aggregated per word with NeMo's `min`.
 
+The optional TDT beam decoder follows NeMo's default sequence-level beam
+search and exposes raw/normalized scores plus token frame/duration metadata.
+See [`docs/tdt-nbest.md`](docs/tdt-nbest.md).
+
 The `parakeet-cli` binary lands at `build/examples/cli/parakeet-cli`.
 
 ---

docs/tdt-nbest.md

@@ -0,0 +1,112 @@
+# TDT N-best decoding
+
+`parakeet.cpp` has an opt-in offline beam decoder for TDT checkpoints. It
+returns multiple complete transcript hypotheses without changing the existing
+greedy decoder.
+
+## Algorithm
+
+`pk::tdt_beam_search` follows NVIDIA NeMo 2.7.3's
+[`BeamTDTInfer.default_beam_search`](https://github.com/NVIDIA-NeMo/NeMo/blob/v2.7.3/nemo/collections/asr/parts/submodules/tdt_beam_decoding.py#L346-L479):
+
+1. Apply `log_softmax` independently to token and duration logits.
+2. Expand the highest-scoring non-blank token-duration pairs.
+3. Expand blank only with non-zero durations.
+4. Merge paths with the same token sequence and final frame using `logaddexp`.
+5. Sort by `score / (emitted_tokens + 1)` by default. The extra item is NeMo's
+   leading blank/SOS sentinel.
+
+The C++ loop additionally rejects non-finite log probabilities and a
+zero-duration expansion that does not strictly reduce the accumulated score.
+This prevents malformed logits or floating-point saturation from creating a
+non-progress loop without truncating valid hypotheses.
+
+There is one intentional edge-case difference from NeMo's reference loop.
+When the duration beam contains only duration 0, blank cannot use that
+duration. NeMo substitutes the smallest positive duration; this implementation
+substitutes the highest-scoring positive duration so the legal blank expansion
+preserves model score ordering. A model-independent regression covers this
+case for `beam_size=1`.
+
+The implementation is limited to the default sequence-level TDT beam search.
+It does not add mAES/mALSD, an external language model, batched beam search, or
+streaming beam search. [NeMo labels this strategy experimental and recommends
+`malsd_batch`](https://github.com/NVIDIA-NeMo/NeMo/blob/v2.7.3/nemo/collections/asr/parts/submodules/rnnt_decoding.py#L513-L533);
+this first implementation deliberately prioritizes a small, directly
+comparable reference algorithm over a second decoder design.
+
+## CLI
+
+```sh
+parakeet-cli transcribe \
+  --model parakeet-tdt-0.6b-v3-q8_0.gguf \
+  --input audio.wav \
+  --decoder tdt \
+  --beam-size 4 \
+  --nbest 4
+```
+
+The command emits JSON:
+
+```json
+{
+  "beam_size": 4,
+  "score_norm": true,
+  "frame_sec": 0.08,
+  "hypotheses": [
+    {
+      "text": "A complete transcript.",
+      "score": -0.941021,
+      "normalized_score": -0.020912,
+      "tokens": [
+        {
+          "id": 499,
+          "frame": 4,
+          "t": 0.32,
+          "duration_frames": 1,
+          "duration": 0.08
+        }
+      ]
+    }
+  ]
+}
+```
+
+Use `--no-score-norm` to rank by raw accumulated score. `--nbest` defaults to
+the beam size.
+
+## C API
+
+The additive entry points are:
+
+```c
+parakeet_capi_transcribe_path_nbest_json(...);
+parakeet_capi_transcribe_pcm_nbest_json(...);
+```
+
+Both reuse a load-once `parakeet_ctx`, accept an optional language prompt, and
+return the same JSON document as the CLI. Existing symbols and ABI version 5
+are unchanged.
+
+## Parity test
+
+Generate an authoritative NeMo baseline:
+
+```sh
+.venv/bin/python scripts/gen_nemo_baseline.py \
+  --model nvidia/parakeet-tdt_ctc-110m \
+  --audio tests/fixtures/speech.wav \
+  --tdt-beam-size 4 \
+  --output /tmp/tdt_nbest_baseline.gguf
+```
+
+Then compare token IDs, token end frames, rank order, and raw scores:
+
+```sh
+PARAKEET_TEST_GGUF=/tmp/pk110m.gguf \
+PARAKEET_TEST_TDT_NBEST_BASELINE=/tmp/tdt_nbest_baseline.gguf \
+ctest --test-dir build --output-on-failure -R '^test_tdt_beam$'
+```
+
+The pull-request closed-loop job generates this baseline directly with NeMo
+and runs the comparison.

examples/cli/main.cpp

@@ -201,6 +201,9 @@ static int cmd_transcribe(int argc, char** argv) {
     bool stream = false;
     bool timestamps = false;
     bool json = false;
+    bool score_norm = true;
+    int beam_size = 0;
+    int nbest = 0;
     int threads = 0;  // 0 == unset -> use the persistent-backend default
     for (int i = 0; i < argc; ++i) {
         if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) {
@@ -219,20 +222,32 @@ static int cmd_transcribe(int argc, char** argv) {
             threads = std::atoi(argv[++i]);
         } else if (std::strcmp(argv[i], "--json") == 0) {
             json = true;
+        } else if (std::strcmp(argv[i], "--beam-size") == 0 && i + 1 < argc) {
+            beam_size = std::atoi(argv[++i]);
+        } else if (std::strcmp(argv[i], "--nbest") == 0 && i + 1 < argc) {
+            nbest = std::atoi(argv[++i]);
+        } else if (std::strcmp(argv[i], "--no-score-norm") == 0) {
+            score_norm = false;
         }
     }
     if (model.empty() || input.empty()) {
         std::fprintf(stderr,
             "usage: parakeet-cli transcribe --model <m.gguf> --input <wav|-> "
             "[--decoder ctc|tdt] [--lang <locale>] [--stream] [--timestamps] "
-            "[--threads N] [--json]\n");
+            "[--threads N] [--json] "
+            "[--beam-size N [--nbest N] [--no-score-norm]]\n");
         return 2;
     }
     // Apply the thread override (offline + streaming graph compute). When unset
     // the persistent-backend default (kDefaultThreads) is used.
     if (threads > 0) pk::set_num_threads(threads);
 
     if (stream) {
+        if (beam_size != 0 || nbest != 0) {
+            std::fprintf(stderr,
+                "parakeet-cli: --beam-size/--nbest are offline TDT only\n");
+            return 2;
+        }
         if (!decoder_str.empty()) {
             std::fprintf(stderr,
                 "parakeet-cli: --stream is RNN-T only; --decoder is ignored\n");
@@ -262,6 +277,89 @@ static int cmd_transcribe(int argc, char** argv) {
         }
     }
 
+    if (nbest != 0 && beam_size == 0) {
+        std::fprintf(stderr,
+            "parakeet-cli: --nbest requires --beam-size\n");
+        return 2;
+    }
+    if (!score_norm && beam_size == 0) {
+        std::fprintf(stderr,
+            "parakeet-cli: --no-score-norm requires --beam-size\n");
+        return 2;
+    }
+    if (beam_size != 0) {
+        if (beam_size < 1) {
+            std::fprintf(stderr,
+                "parakeet-cli: --beam-size must be at least 1\n");
+            return 2;
+        }
+        if (nbest == 0) nbest = beam_size;
+        if (nbest < 1 || nbest > beam_size) {
+            std::fprintf(stderr,
+                "parakeet-cli: require beam-size >= nbest >= 1\n");
+            return 2;
+        }
+        if (decoder_str == "ctc") {
+            std::fprintf(stderr,
+                "parakeet-cli: TDT N-best cannot use --decoder ctc\n");
+            return 2;
+        }
+        if (timestamps) {
+            std::fprintf(stderr,
+                "parakeet-cli: --timestamps is redundant with N-best JSON\n");
+            return 2;
+        }
+
+        if (is_stdin_input(input)) {
+            pk::Audio audio;
+            if (!load_audio_arg_16k_mono(input, audio)) {
+                std::fprintf(stderr, "parakeet-cli: failed to load audio stdin\n");
+                return 1;
+            }
+            try {
+                std::unique_ptr<pk::Model> m = pk::Model::load(model);
+                if (!m) {
+                    std::fprintf(stderr,
+                        "parakeet-cli: failed to load model %s\n", model.c_str());
+                    return 1;
+                }
+                std::vector<pk::NBestTranscription> hypotheses =
+                    m->transcribe_pcm_nbest(
+                        audio.samples, audio.sample_rate,
+                        beam_size, nbest, score_norm, lang);
+                std::string output = pk::nbest_transcriptions_to_json(
+                    hypotheses, beam_size, score_norm, model_frame_sec(*m));
+                std::printf("%s\n", output.c_str());
+            } catch (const std::exception& e) {
+                std::fprintf(stderr,
+                    "parakeet-cli: N-best transcription failed: %s\n", e.what());
+                return 1;
+            }
+            return 0;
+        }
+
+        parakeet_ctx* ctx = parakeet_capi_load(model.c_str());
+        if (!ctx) {
+            std::fprintf(stderr,
+                "parakeet-cli: failed to load model %s\n", model.c_str());
+            return 1;
+        }
+        char* output = parakeet_capi_transcribe_path_nbest_json(
+            ctx, input.c_str(), beam_size, nbest, score_norm ? 1 : 0,
+            lang.empty() ? nullptr : lang.c_str());
+        if (!output) {
+            std::fprintf(stderr,
+                "parakeet-cli: N-best transcription failed: %s\n",
+                parakeet_capi_last_error(ctx));
+            parakeet_capi_free(ctx);
+            return 1;
+        }
+        std::printf("%s\n", output);
+        parakeet_capi_free_string(output);
+        parakeet_capi_free(ctx);
+        return 0;
+    }
+
     // --json: emit the C-API JSON document (text + word/token timestamps + conf).
     if (json) {
         if (is_stdin_input(input)) {
@@ -1270,7 +1368,8 @@ int main(int argc, char** argv) {
         "  parakeet-cli info <model.gguf>\n"
         "  parakeet-cli transcribe --model <model.gguf> --input <wav|-> "
         "[--decoder ctc|tdt] [--lang <locale>] [--stream] [--timestamps] "
-        "[--threads N] [--json]\n"
+        "[--threads N] [--json] "
+        "[--beam-size N [--nbest N] [--no-score-norm]]\n"
         "  parakeet-cli quantize <in.gguf> <out.gguf> "
         "<q4_0|q5_0|q8_0|q4_k|q5_k|q6_k>\n"
         "  parakeet-cli bench --model <model.gguf> --manifest <file> "

include/parakeet_capi.h

@@ -159,6 +159,27 @@ char* parakeet_capi_transcribe_pcm_batch_json_lang(parakeet_ctx* ctx,
                                                    int sample_rate, int decoder,
                                                    const char* target_lang);
 
+// Offline TDT beam search returning ranked hypotheses as JSON. These are
+// additive, TDT-only entry points; they do not change the existing greedy
+// transcription functions or ABI version. `beam_size >= nbest >= 1`.
+// `score_norm != 0` matches NeMo's default score/sequence-length ranking.
+// `target_lang` has the same semantics as the other *_lang functions.
+//
+// JSON shape:
+//   {"beam_size":4,"score_norm":true,"frame_sec":0.080000,
+//    "hypotheses":[
+//      {"text":"...","score":-12.3,"normalized_score":-0.45,
+//       "tokens":[{"id":123,"frame":7,"t":0.560,
+//                  "duration_frames":2,"duration":0.160}, ...]}
+//    ]}
+char* parakeet_capi_transcribe_path_nbest_json(
+    parakeet_ctx* ctx, const char* wav_path,
+    int beam_size, int nbest, int score_norm, const char* target_lang);
+
+char* parakeet_capi_transcribe_pcm_nbest_json(
+    parakeet_ctx* ctx, const float* samples, int n_samples, int sample_rate,
+    int beam_size, int nbest, int score_norm, const char* target_lang);
+
 // ---------------------------------------------------------------------------
 // Streaming API (cache-aware streaming RNN-T, e.g. the EOU model
 // nvidia/parakeet_realtime_eou_120m-v1). The stream session buffers incoming