diff --git a/parakeet-cpp/CMakeLists.txt b/parakeet-cpp/CMakeLists.txt index 5a0bca317b3..989016507e3 100644 --- a/parakeet-cpp/CMakeLists.txt +++ b/parakeet-cpp/CMakeLists.txt @@ -488,6 +488,9 @@ if (PARAKEET_BUILD_TESTS) set(_qvp_ctc_ref "${PARAKEET_TEST_REF_DIR}/ctc-ref") set(_qvp_tdt_ref "${PARAKEET_TEST_REF_DIR}/tdt-ref") set(_qvp_sf_ref "${PARAKEET_TEST_REF_DIR}/sortformer-ref") + set(_qvp_rnnt_q8_gguf "${PARAKEET_TEST_MODEL_DIR}/stt_ka_fastconformer_hybrid_large_pc.q8_0.gguf") + set(_qvp_rnnt_wav "${PARAKEET_TEST_AUDIO_DIR}/rnnt-ka-16k.wav") + set(_qvp_rnnt_ref "${PARAKEET_TEST_REF_DIR}/rnnt-ref") add_executable(test-mel test/test_mel.cpp @@ -747,6 +750,30 @@ if (PARAKEET_BUILD_TESTS) ARGS "${_qvp_tdt_q8_gguf}" "${_qvp_fr_multipiece_wav}" REQUIRES "${_qvp_tdt_q8_gguf}" "${_qvp_fr_multipiece_wav}") + # Plain RNN-T greedy decoder parity: [ref-dir]. The + # rnnt GGUF has no public checkpoint in the fixture set (converted from + # a hybrid NeMo checkpoint, e.g. the Georgian model), so the test is + # registered with REQUIRES: it shows in `ctest -N` and stays DISABLED + # until the fixture files are dropped in place. + add_executable(test-rnnt-decoder-parity + test/test_rnnt_decoder_parity.cpp + src/parakeet_ctc.cpp + src/parakeet_engine.cpp + src/parakeet_log.cpp + src/parakeet_tdt.cpp + src/parakeet_eou.cpp + src/parakeet_sortformer.cpp + src/mel_preprocess.cpp + src/sentencepiece_bpe.cpp + src/energy_vad.cpp) + target_link_libraries(test-rnnt-decoder-parity PRIVATE ggml parakeet-backend-defs) + target_include_directories(test-rnnt-decoder-parity PRIVATE include src ggml/include) + parakeet_apply_backend_defs(test-rnnt-decoder-parity) + parakeet_register_test(test-rnnt-decoder-parity + LABEL "fixture" + ARGS "${_qvp_rnnt_q8_gguf}" "${_qvp_rnnt_wav}" "${_qvp_rnnt_ref}" + REQUIRES "${_qvp_rnnt_q8_gguf}" "${_qvp_rnnt_wav}" "${_qvp_rnnt_ref}/token_ids.npy") + add_executable(test-sortformer-parity test/test_sortformer_parity.cpp src/parakeet_ctc.cpp diff --git a/parakeet-cpp/PROGRESS.md b/parakeet-cpp/PROGRESS.md index 235093e2c1e..e3b10c5d775 100644 --- a/parakeet-cpp/PROGRESS.md +++ b/parakeet-cpp/PROGRESS.md @@ -3470,3 +3470,91 @@ doesn't flake the test. binary. Surfacing it through downstream addon wrappers (e.g. `transcription-parakeet`'s `runStreaming()` JS API) requires separate plumbing work on those wrappers — not in this phase. + +## Phase 18 — plain RNN-T (Transducer) head _(done)_ + +### 18.0 — scope and design decision + +Duration-less Transducer checkpoints had no working path: the converter +mis-tagged them `tdt` and crashed on the absent +`model_defaults.tdt_durations`; the TDT loader assumes duration logits +(`V+1+num_durations` joint rows); the EOU decoder is itself a plain +RNN-T but hard-wires ``/`` semantics (and its `eou_id` +fallback `vocab_size-2` would alias real BPE ids on a model without +those tokens). Target checkpoint family: the RNN-T branch of hybrid +`EncDecHybridRNNTCTCBPEModel` models, concretely +`nvidia/stt_ka_fastconformer_hybrid_large_pc` (Georgian). + +Design: plain RNN-T is structurally EOU minus the special tokens, so +the head reuses the EOU predictor/joint runtime (`EouRuntimeWeights`, +`eou_prepare_runtime`, `eou_decode_window`, `eou_greedy_decode`) with a +new `EouDecodeOptions.disable_special_tokens` flag instead of adding a +fourth decoder implementation. In plain mode the greedy inner loop +breaks only on the transducer blank — matching NeMo greedy RNN-T — +and `eou_id`/`eob_id` are forced to `-1` (never valid ids). + +### 18.1 — converter + +`--head {auto,ctc,tdt,rnnt,eou,sortformer}` override; +`detect_model_type()` now routes RNN-T checkpoints with neither TDT +durations nor an `` label to `rnnt` instead of crashing as `tdt`. +Emits `parakeet.rnnt.{vocab_size,blank_id,pred_hidden,pred_rnn_layers, +joint_hidden,max_symbols_per_step}` (blank_as_pad: `blank_id == +vocab_size`) + `rnnt.predict.*` / `rnnt.joint.*` tensors; the hybrid's +CTC aux head (`ctc_decoder.*`) is not exported. Guards: joint output +rows must equal `vocab+1` (a TDT-shaped joint forced through `--head +rnnt` fails at convert time, not at decode time), and +`joint.num_classes` must equal `decoder.vocab_size`. + +### 18.2 — loader + Engine dispatch + +`ParakeetModelType::RNNT`; `rnnt.*` tensors load into the shared EOU +weight slots. Decode arms in `Engine::transcribe_samples`, both +streaming paths (`transcribe_samples_stream`, `StreamSession` +`process_window` — mirrors the EOU arms with the plain flag and +persistent decode state), `stream_start` state init, and the CLI. +`is_transcription_model()` includes RNNT; `model_type()` returns +`"rnnt"`. Plain RNNT has no native end-pointing signal, so streaming +sessions take the same opt-in RMS energy VAD as CTC/TDT. + +### 18.3 — parity validation + +`scripts/dump-rnnt-reference.py` (forces +`change_decoding_strategy(decoder_type="rnnt")` on the hybrid, dumps +greedy `token_ids.npy` + transcript + optional `encoder_out.npy`) and +`test/test_rnnt_decoder_parity.cpp` (CPU vs GPU-offloaded-encoder run, +plus bit-exact token-id comparison against the NeMo dump; int32 C-order +`.npy` enforced; wav sample rate validated against the model). + +Measured 2026-06-09 on `nvidia/stt_ka_fastconformer_hybrid_large_pc`: +C++ greedy decode reproduces NeMo greedy token ids **75/75 bit-for-bit** +on a FLEURS `ka` clip; CPU == Metal at f16; Metal == NeMo at q8_0; +FLEURS `ka` WER within **0.31 %** of NeMo. The ctest entry +(`test-rnnt-decoder-parity`, label `fixture`) registers via `REQUIRES` +and stays DISABLED until the external GGUF + wav + ref dump are placed +under the fixture dirs (no public checkpoint for this head). + +### 18.4 — files touched + +- `scripts/convert-nemo-to-gguf.py` — `--head`, rnnt detection + + metadata + tensors + shape guards. +- `scripts/dump-rnnt-reference.py` (new) — NeMo greedy reference dump. +- `src/parakeet_ctc.{h,cpp}` — `ParakeetModelType::RNNT`, rnnt GGUF + loader arm, model summary. +- `src/parakeet_eou.{h,cpp}` — `disable_special_tokens` plain mode; + RNNT accepted by `eou_prepare_runtime` with `-1` sentinel ids. +- `src/parakeet_engine.cpp` — offline + Mode 2 + Mode 3 dispatch, + `stream_start` init, energy-VAD comment scope. +- `src/main.cpp` — CLI decode arm + usage text. +- `test/test_rnnt_decoder_parity.cpp` (new) + `CMakeLists.txt` + fixture vars and `parakeet_register_test` entry. + +### 18.5 — follow-ups + +- Streaming-window decode (Mode 2/3) has no dedicated parity harness + yet; it reuses the EOU state machinery verified in §12 but a + windowed-vs-offline token-id cross-check on a long clip would close + the gap. +- Decode is scalar-CPU (like EOU). If the fused TDT Metal decoder + (§15) ever generalises to `num_durations == 0`, plain RNNT can ride + the same graph path. diff --git a/parakeet-cpp/README.md b/parakeet-cpp/README.md index 0c694f40dcc..2e359245dcc 100644 --- a/parakeet-cpp/README.md +++ b/parakeet-cpp/README.md @@ -1,6 +1,6 @@ # parakeet.cpp -**Parakeet** (NVIDIA FastConformer ASR family, CC-BY-4.0) ported to [`ggml`](https://github.com/ggml-org/ggml). Pure C++ inference on **CPU** and **GPU** (Metal / Vulkan / OpenCL); no Python, PyTorch, or onnxruntime at runtime. One **`parakeet::Engine`** loads **CTC**, **TDT**, **EOU**, or **Sortformer** GGUFs and dispatches by metadata. +**Parakeet** (NVIDIA FastConformer ASR family, CC-BY-4.0) ported to [`ggml`](https://github.com/ggml-org/ggml). Pure C++ inference on **CPU** and **GPU** (Metal / Vulkan / OpenCL); no Python, PyTorch, or onnxruntime at runtime. One **`parakeet::Engine`** loads **CTC**, **TDT**, **RNNT**, **EOU**, or **Sortformer** GGUFs and dispatches by metadata. ## Supported checkpoints @@ -17,22 +17,24 @@ Encoder topology is selected from GGUF metadata (`conv_norm_type`, causal subsampling, chunked-limited attention, etc.), so EOU shares the same C++ graph path as CTC/TDT where weights allow. +Plain **RNN-T** (`parakeet.model.type = "rnnt"`) covers duration-less Transducer heads — e.g. the RNN-T branch of a hybrid `EncDecHybridRNNTCTCBPEModel` checkpoint (`--head rnnt` in the converter; the hybrid's CTC aux head is ignored). There is no public checkpoint in the fixture set; the head is verified against NeMo greedy decoding on `nvidia/stt_ka_fastconformer_hybrid_large_pc` (Georgian) via `test-rnnt-decoder-parity` + `scripts/dump-rnnt-reference.py`. The decoder reuses the EOU predictor/joint runtime (scalar CPU) with special-token handling disabled. + ## API overview | Surface | Role | |---------|------| -| `Engine::transcribe` | One-shot wav → text (CTC / TDT / EOU) or segments (Sortformer) | +| `Engine::transcribe` | One-shot wav → text (CTC / TDT / RNNT / EOU) or segments (Sortformer) | | `Engine::transcribe_stream` | Mode 2: full encode once, stream segments | | `Engine::stream_start` → `StreamSession` | Mode 3: live duplex / cache-aware chunks | | `Engine::diarize` / `diarize_start` | Sortformer offline / live streaming (v1: sliding-history; v2.1: speaker-cache / AOSC) | | `transcribe_with_speakers` | Sortformer + ASR → attributed transcript | -EOU streaming segments expose `is_eou_boundary`. **`StreamEvent`** (optional callbacks) covers end-of-turn (EOU) and VAD-style signals (Sortformer threshold, optional energy VAD on CTC/TDT). **`Engine::backend_device`** / **`backend_name`** reflect the backend actually used after the load-time cascade. +EOU streaming segments expose `is_eou_boundary`. **`StreamEvent`** (optional callbacks) covers end-of-turn (EOU) and VAD-style signals (Sortformer threshold, optional energy VAD on CTC/TDT/RNNT). **`Engine::backend_device`** / **`backend_name`** reflect the backend actually used after the load-time cascade. ## Pipeline ``` -wav → log-mel → FastConformer encoder → CTC / TDT / EOU / Sortformer decoder +wav → log-mel → FastConformer encoder → CTC / TDT / RNNT / EOU / Sortformer decoder ``` Each GGUF bundles weights, mel filterbank, and tokenizer as needed. @@ -107,6 +109,8 @@ python scripts/convert-nemo-to-gguf.py \ **Important:** for non-default checkpoints set **`--hf-repo`** (e.g. `nvidia/parakeet-tdt-0.6b-v3`) — the script otherwise defaults to the CTC repo and may download the wrong weights. Use `scripts/download-all-models.sh` to prefetch `.nemo` files. +The head is auto-detected from `cfg['target']`; **`--head {ctc,tdt,rnnt,eou,sortformer}`** overrides it. Hybrid transducer+CTC checkpoints (`EncDecHybridRNNTCTCBPEModel`) export their plain RNN-T branch as `--head rnnt` (also the auto-detect result when the checkpoint has neither TDT durations nor an `` token); the hybrid's CTC aux head is ignored. + Default **`--quant`** is **`q8_0`**. Use **`f16`** for parity-calibrated harnesses (noise from q8 swamps NeMo FP32 references). ### Quantization tiers (CTC 0.6B, M4 Air CPU) @@ -150,7 +154,7 @@ CMake builds the main binary as target **`parakeet-cli`** with **`OUTPUT_NAME pa **Synopsis:** `parakeet --model <.gguf> (--wav <.wav> | --pcm-in <.raw>) [options]` -The GGUF picks the engine (CTC / TDT / EOU transcription vs Sortformer diarization). Optional **`--diarization-model `** adds speaker labels when **`--model`** is a CTC/TDT GGUF (“who said what”). +The GGUF picks the engine (CTC / TDT / RNNT / EOU transcription vs Sortformer diarization). Optional **`--diarization-model `** adds speaker labels when **`--model`** is a CTC/TDT/RNNT GGUF (“who said what”). | Topic | Flags | |------|--------| @@ -259,6 +263,8 @@ python scripts/dump-ctc-reference.py --wav test/samples/jfk.wav python scripts/dump-tdt-reference.py --wav test/samples/jfk.wav python scripts/dump-eou-reference.py --wav test/samples/jfk.wav python scripts/dump-sortformer-reference.py --wav test/samples/diarization-sample-16k.wav +# rnnt: no fixture wav in-repo; point --wav at your own 16 kHz clip +python scripts/dump-rnnt-reference.py --nemo-model --wav cmake -S . -B build -DCMAKE_BUILD_TYPE=Release cmake --build build -j diff --git a/parakeet-cpp/include/parakeet/attributed.h b/parakeet-cpp/include/parakeet/attributed.h index 039e53d8487..93eec2e8097 100644 --- a/parakeet-cpp/include/parakeet/attributed.h +++ b/parakeet-cpp/include/parakeet/attributed.h @@ -1,6 +1,6 @@ #pragma once -// Speaker-attributed transcription: Sortformer segments + ASR text per slice (CTC/TDT/EOU). +// Speaker-attributed transcription: Sortformer segments + ASR text per slice (CTC/TDT/RNNT/EOU). #include "export.h" #include "engine.h" diff --git a/parakeet-cpp/include/parakeet/engine.h b/parakeet-cpp/include/parakeet/engine.h index 95cbec1f9c3..ca32ad74d0a 100644 --- a/parakeet-cpp/include/parakeet/engine.h +++ b/parakeet-cpp/include/parakeet/engine.h @@ -3,7 +3,7 @@ // Loaded GGUF inference: transcribe, stream, diarize, and backend metadata behind one Engine class. // // Loads weights once; subsequent calls pay mel + encoder + decode only. Model kind (CTC, TDT, -// EOU, Sortformer) comes from GGUF metadata. +// RNNT, EOU, Sortformer) comes from GGUF metadata. // // Transcription: // - transcribe / transcribe_samples — one-shot wav or PCM to text. @@ -230,7 +230,7 @@ class PARAKEET_API Engine { const EngineOptions & options() const; - // "ctc", "tdt", "eou", or "sortformer", reflecting the + // "ctc", "tdt", "rnnt", "eou", or "sortformer", reflecting the // parakeet.model.type metadata of the loaded GGUF. std::string model_type() const; diff --git a/parakeet-cpp/include/parakeet/parakeet.h b/parakeet-cpp/include/parakeet/parakeet.h index 2f8514a5f32..d9b1860fb93 100644 --- a/parakeet-cpp/include/parakeet/parakeet.h +++ b/parakeet-cpp/include/parakeet/parakeet.h @@ -6,8 +6,8 @@ // - parakeet_cli_main C entry point // - parakeet_log_set host log sink // - Engine + EngineOptions / EngineResult -// (CTC, TDT, EOU, Sortformer behind one -// class) +// (CTC, TDT, RNNT, EOU, Sortformer behind +// one class) // - StreamingOptions / StreamingSegment / // StreamSession + cross-engine // StreamEvent + VadState + diff --git a/parakeet-cpp/include/parakeet/streaming.h b/parakeet-cpp/include/parakeet/streaming.h index da4c24adb68..1001ca4dfb4 100644 --- a/parakeet-cpp/include/parakeet/streaming.h +++ b/parakeet-cpp/include/parakeet/streaming.h @@ -17,7 +17,7 @@ namespace parakeet { // Optional StreamEvent callback: VadStateChanged and EndOfTurn alongside segment text. // // EOU models emit EndOfTurn when `` fires. Sortformer emits VadStateChanged from -// speaker_probs vs threshold. CTC/TDT can use optional RMS EnergyVad when enabled. +// speaker_probs vs threshold. CTC/TDT/RNNT can use optional RMS EnergyVad when enabled. enum class VadState : int { Unknown = 0, @@ -67,12 +67,12 @@ struct StreamingOptions { // Optional; nullptr disables StreamEvent delivery (segment-only streaming). StreamEventCallback on_event = nullptr; - // Energy-VAD fallback. When true, CTC / TDT sessions will compute a - // simple RMS-thresholded VAD over the input PCM and fire + // Energy-VAD fallback. When true, CTC / TDT / RNNT sessions will compute + // a simple RMS-thresholded VAD over the input PCM and fire // `StreamEventType::VadStateChanged` events on transitions. Always-on // for sessions whose underlying engine (EOU, Sortformer) has its own // native VAD source -- those engines' events take priority. Default - // off; opt-in for CTC/TDT consumers that want VadState events. + // off; opt-in for CTC/TDT/RNNT consumers that want VadState events. bool enable_energy_vad = false; // Energy-VAD knobs (dB-scale; applies only when enable_energy_vad). @@ -109,7 +109,7 @@ struct StreamingSegment { // segment whose token list is empty (defensive default). bool starts_word = true; - // EOU-only: true when this segment ends on ``. For CTC/TDT use StreamEvent + // EOU-only: true when this segment ends on ``. For CTC/TDT/RNNT use StreamEvent // EndOfTurn via `on_event` instead; those engines leave this flag false here. bool is_eou_boundary = false; float eot_confidence = 0.0f; diff --git a/parakeet-cpp/scripts/convert-nemo-to-gguf.py b/parakeet-cpp/scripts/convert-nemo-to-gguf.py index aed3a2314e1..76f59b2b2e2 100644 --- a/parakeet-cpp/scripts/convert-nemo-to-gguf.py +++ b/parakeet-cpp/scripts/convert-nemo-to-gguf.py @@ -15,10 +15,16 @@ cache-aware streaming, end-of- utterance token detection; parakeet_realtime_eou_120m-v1) + - ``EncDecHybridRNNTCTCBPEModel`` / ``EncDecRNNTBPEModel`` + (no TDT durations, no ```` token) + -> plain RNN-T (Transducer) head; + for hybrids only the Transducer + branch is exported (the CTC aux + head ``ctc_decoder.*`` is ignored) - ``EncDecDiarLabelModel`` -> Sortformer (diar_sortformer_4spk-v1, diar_streaming_sortformer_4spk-v2) -The FastConformer encoder topology is shared across all four flavours; only +The FastConformer encoder topology is shared across all five flavours; only the decoder / head tensors + metadata differ. EOU additionally swaps the conv module's BatchNorm for a LayerNorm and carries cache-aware streaming hyperparameters (att_context_size, subsampling-output cache lookback, and the @@ -34,7 +40,7 @@ Metadata: general.architecture = "parakeet-ctc" (kept for GGUF compat) general.name = "" - parakeet.model.type = "ctc", "tdt", "eou", or "sortformer" + parakeet.model.type = "ctc", "tdt", "rnnt", "eou", or "sortformer" parakeet.encoder.* (hyperparameters, incl. use_bias, xscaling, conv_norm_type, att_context_size, causal_downsampling, conv_context_size) @@ -42,6 +48,9 @@ parakeet.ctc.* (vocab_size, blank_id) [CTC only] parakeet.tdt.* (predictor + joint hyperparameters + durations) [TDT only] + parakeet.rnnt.* (vocab_size, blank_id, pred_hidden, + pred_rnn_layers, joint_hidden, + max_symbols_per_step) [RNNT only] parakeet.eou.* (vocab_size, blank_id, eou_id, eob_id, pred_hidden, pred_rnn_layers, joint_hidden, encoder_chunk_mel_frames, @@ -64,6 +73,10 @@ tdt.predict.lstm.{l}.{w_ih,w_hh,b_ih,b_hh} [TDT only] tdt.joint.{enc,pred}.{weight,bias} [TDT only] tdt.joint.out.{weight,bias} [TDT only] + rnnt.predict.embed.weight [RNNT only] + rnnt.predict.lstm.{l}.{w_ih,w_hh,b_ih,b_hh} [RNNT only] + rnnt.joint.{enc,pred}.{weight,bias} [RNNT only] + rnnt.joint.out.{weight,bias} [RNNT only] eou.predict.embed.weight [EOU only] eou.predict.lstm.0.{w_ih,w_hh,b_ih,b_hh} [EOU only] eou.joint.{enc,pred}.{weight,bias} [EOU only] @@ -117,6 +130,12 @@ def parse_args() -> argparse.Namespace: "pass --quant f16 for the bit-equal floating-point baseline.") p.add_argument("--hf-repo", default="nvidia/parakeet-ctc-0.6b", help="HF model id to download from if --ckpt is missing.") + p.add_argument("--head", choices=["auto", "ctc", "tdt", "rnnt", "eou", "sortformer"], + default="auto", + help="Override the auto-detected head. 'auto' (default) infers from " + "cfg['target']. Use 'rnnt' to force the plain RNN-T (Transducer) " + "head of a hybrid EncDecHybridRNNTCTCBPEModel checkpoint; its CTC " + "aux head (ctc_decoder.*) is ignored.") return p.parse_args() @@ -168,7 +187,9 @@ def load_nemo(ckpt: Path): return cfg, sd, tok_bytes -def detect_model_type(cfg: dict) -> str: +def detect_model_type(cfg: dict, head: str = "auto") -> str: + if head != "auto": + return head target = cfg.get("target", "") if "Sortformer" in target or "sortformer_modules" in cfg: return "sortformer" @@ -182,7 +203,11 @@ def detect_model_type(cfg: dict) -> str: has_eou = any(str(lbl) == "" for lbl in labels) if has_eou: return "eou" - return "tdt" + # RNN-T with neither TDT durations nor an token: a plain + # Transducer head (e.g. the RNN-T branch of a hybrid + # EncDecHybridRNNTCTCBPEModel). Previously mis-tagged "tdt", which + # crashed below on the absent model_defaults.tdt_durations. + return "rnnt" return "ctc" @@ -216,8 +241,9 @@ def detect_sortformer_variant(ckpt: Path) -> str: return "" -def write_gguf(out: Path, ckpt: Path, cfg: dict, sd: dict, tok_bytes: bytes, quant: str): - model_type = detect_model_type(cfg) +def write_gguf(out: Path, ckpt: Path, cfg: dict, sd: dict, tok_bytes: bytes, quant: str, + head: str = "auto"): + model_type = detect_model_type(cfg, head) enc = cfg["encoder"] pre = cfg["preprocessor"] @@ -252,6 +278,7 @@ def write_gguf(out: Path, ckpt: Path, cfg: dict, sd: dict, tok_bytes: bytes, qua model_name = { "ctc": f"parakeet-ctc-{d_model}-{n_layers}l", "tdt": f"parakeet-tdt-{d_model}-{n_layers}l", + "rnnt": f"parakeet-rnnt-{d_model}-{n_layers}l", "eou": f"parakeet-eou-{d_model}-{n_layers}l", "sortformer": f"sortformer-{d_model}-{n_layers}l", }[model_type] @@ -354,6 +381,45 @@ def write_gguf(out: Path, ckpt: Path, cfg: dict, sd: dict, tok_bytes: bytes, qua variant = detect_sortformer_variant(ckpt) if variant: writer.add_string("parakeet.model_variant", variant) + elif model_type == "rnnt": + # Plain RNN-T (Transducer) head. Structurally a TDT head minus the + # duration outputs: joint.out is (vocab+1, joint_hidden) where the +1 + # is blank only (no duration logits). Predictor + joint tensor keys are + # identical to TDT/EOU; the hybrid's CTC aux head (ctc_decoder.*) is + # ignored. Cache-aware streaming params (att_context_size, conv_norm_type) + # ride in the shared parakeet.encoder.* metadata already emitted above. + if "prednet" not in dec or "joint" not in cfg: + raise RuntimeError( + "rnnt head requires a Transducer checkpoint " + "(decoder.prednet / joint config missing -- is this a CTC model?)") + pred_hidden = int(dec["prednet"]["pred_hidden"]) + pred_rnn_layers = int(dec["prednet"]["pred_rnn_layers"]) + joint_hidden = int(cfg["joint"]["jointnet"]["joint_hidden"]) + pred_vocab_size = int(dec["vocab_size"]) # label vocab (no blank) + joint_num_classes = int(cfg["joint"]["num_classes"]) # label vocab, blank excluded (RNNTJoint adds +1) + blank_id = joint_num_classes # blank_as_pad at index vocab_size + if joint_num_classes != pred_vocab_size: + raise RuntimeError( + f"rnnt: joint.num_classes ({joint_num_classes}) != decoder.vocab_size " + f"({pred_vocab_size}); blank_id placement would be wrong") + greedy_cfg = (cfg.get("decoding") or {}).get("greedy") or {} + max_symbols_cfg = greedy_cfg.get("max_symbols", + greedy_cfg.get("max_symbols_per_step")) + if not max_symbols_cfg: + # NeMo treats an unset/None max_symbols as "no cap"; the C++ greedy + # loop needs a finite cap, so bake in NeMo's usual default of 10. + print("[convert] rnnt: decoding.greedy.max_symbols unset (uncapped in " + "NeMo); capping at 10 in the GGUF", file=sys.stderr) + max_symbols = 10 + else: + max_symbols = int(max_symbols_cfg) + + writer.add_uint32("parakeet.rnnt.vocab_size", pred_vocab_size) + writer.add_uint32("parakeet.rnnt.blank_id", blank_id) + writer.add_uint32("parakeet.rnnt.pred_hidden", pred_hidden) + writer.add_uint32("parakeet.rnnt.pred_rnn_layers", pred_rnn_layers) + writer.add_uint32("parakeet.rnnt.joint_hidden", joint_hidden) + writer.add_uint32("parakeet.rnnt.max_symbols_per_step", max_symbols) else: pred_hidden = int(dec["prednet"]["pred_hidden"]) pred_rnn_layers = int(dec["prednet"]["pred_rnn_layers"]) @@ -583,6 +649,36 @@ def try_bias(name: str, key: str): sd["sortformer_modules.single_hidden_to_spks.weight"]) add_f32("sortformer.head.single_hidden_to_spks.bias", sd["sortformer_modules.single_hidden_to_spks.bias"]) + elif model_type == "rnnt": + add_2d ("rnnt.predict.embed.weight", sd["decoder.prediction.embed.weight"]) + + pred_rnn_layers = int(cfg["decoder"]["prednet"]["pred_rnn_layers"]) + for l in range(pred_rnn_layers): + add_2d (f"rnnt.predict.lstm.{l}.w_ih", + sd[f"decoder.prediction.dec_rnn.lstm.weight_ih_l{l}"]) + add_2d (f"rnnt.predict.lstm.{l}.w_hh", + sd[f"decoder.prediction.dec_rnn.lstm.weight_hh_l{l}"]) + add_f32(f"rnnt.predict.lstm.{l}.b_ih", + sd[f"decoder.prediction.dec_rnn.lstm.bias_ih_l{l}"]) + add_f32(f"rnnt.predict.lstm.{l}.b_hh", + sd[f"decoder.prediction.dec_rnn.lstm.bias_hh_l{l}"]) + + add_2d ("rnnt.joint.enc.weight", sd["joint.enc.weight"]) + add_f32("rnnt.joint.enc.bias", sd["joint.enc.bias"]) + add_2d ("rnnt.joint.pred.weight", sd["joint.pred.weight"]) + add_f32("rnnt.joint.pred.bias", sd["joint.pred.bias"]) + # A plain RNN-T joint emits exactly vocab+1 logits (labels + blank). + # A TDT-shaped joint (vocab+1+num_durations rows) reaching this branch + # (e.g. via --head rnnt) must fail here, not decode garbage later. + out_rows = int(sd["joint.joint_net.2.weight"].shape[0]) + vocab_p1 = int(cfg["decoder"]["vocab_size"]) + 1 + if out_rows != vocab_p1: + raise RuntimeError( + f"rnnt: joint output has {out_rows} rows, expected vocab+1 = " + f"{vocab_p1}; duration logits present? (TDT checkpoint -- use " + f"the auto-detected head instead of forcing rnnt)") + add_2d ("rnnt.joint.out.weight", sd["joint.joint_net.2.weight"]) + add_f32("rnnt.joint.out.bias", sd["joint.joint_net.2.bias"]) else: add_2d ("tdt.predict.embed.weight", sd["decoder.prediction.embed.weight"]) @@ -623,6 +719,11 @@ def try_bias(name: str, key: str): f"blank_id={int(cfg['joint']['num_classes'])} eou_id={eou_pos} " f"att_ctx=[{att_ctx_left},{att_ctx_right}] " f"conv_norm={conv_norm_type}") + elif model_type == "rnnt": + vocab_note = (f"rnnt_vocab={int(cfg['decoder']['vocab_size'])} " + f"blank_id={int(cfg['joint']['num_classes'])} " + f"att_ctx=[{att_ctx_left},{att_ctx_right}] " + f"conv_norm={conv_norm_type}") else: vocab_note = f"tdt_vocab={int(cfg['decoder']['vocab_size'])} durations={cfg['model_defaults']['tdt_durations']}" print(f"[convert] wrote {out} ({size_mb:.1f} MiB, type={model_type}, quant={quant}, {vocab_note}, layers={n_layers}, use_bias={use_bias})", file=sys.stderr) @@ -633,7 +734,7 @@ def main(): ckpt = ensure_ckpt(args.ckpt, args.hf_repo) cfg, sd, tok_bytes = load_nemo(ckpt) args.out.parent.mkdir(parents=True, exist_ok=True) - write_gguf(args.out, ckpt, cfg, sd, tok_bytes, args.quant) + write_gguf(args.out, ckpt, cfg, sd, tok_bytes, args.quant, args.head) if __name__ == "__main__": diff --git a/parakeet-cpp/scripts/dump-rnnt-reference.py b/parakeet-cpp/scripts/dump-rnnt-reference.py new file mode 100644 index 00000000000..df31f737710 --- /dev/null +++ b/parakeet-cpp/scripts/dump-rnnt-reference.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Dump NeMo greedy reference for plain RNN-T parity (token IDs + transcript). + +Adapted from dump-tdt-reference.py for a hybrid EncDecHybridRNNTCTCBPEModel +(e.g. nvidia/stt_ka_fastconformer_hybrid_large_pc): forces the RNN-T decoder +(not the CTC aux head) and dumps the greedy token stream so the C++ +test-rnnt-decoder-parity can assert bit-identical greedy decoding. + + / + token_ids.npy (N,) NeMo greedy RNN-T token IDs + transcript.txt NeMo greedy transcript + encoder_out.npy (T_enc, d_model) NeMo encoder output (optional parity) + mel.npy (n_mels, T_mel) post-preprocessor log-mel (optional) + +Greedy decoding is deterministic; the C++ side must reproduce token_ids exactly. +""" + +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +import torch + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--wav", type=Path, required=True, help="Input mono 16 kHz wav") + p.add_argument("--out", type=Path, default=Path("artifacts/rnnt-ref"), help="Output directory for dumps") + p.add_argument("--nemo-model", type=Path, + default=Path("models/stt_ka_fastconformer_hybrid_large_pc.nemo")) + p.add_argument("--device", default="cpu") + p.add_argument("--no-encoder-dump", action="store_true", + help="Skip mel/encoder_out dumps (token parity only)") + return p.parse_args() + + +def main(): + args = parse_args() + args.out.mkdir(parents=True, exist_ok=True) + os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + + import nemo.collections.asr as nemo_asr + + print(f"[rnnt-ref] restoring {args.nemo_model}", file=sys.stderr) + model = nemo_asr.models.ASRModel.restore_from(str(args.nemo_model), map_location=args.device) + model.eval() + model.preprocessor.featurizer.dither = 0.0 + model.preprocessor.featurizer.pad_to = 0 + + # Hybrid RNN-T/CTC: force the RNN-T decoder (keeps the model's default greedy + # decoding strategy -- the configuration NeMo's published WER is measured at). + if hasattr(model, "cur_decoder"): + try: + model.change_decoding_strategy(decoder_type="rnnt") + print("[rnnt-ref] forced decoder_type=rnnt (hybrid model)", file=sys.stderr) + except Exception as e: # noqa: BLE001 + print(f"[rnnt-ref] WARN: change_decoding_strategy failed ({e}); using model default", + file=sys.stderr) + + try: + gd = model.cfg.decoding.greedy + print(f"[rnnt-ref] decoding.greedy.max_symbols={gd.get('max_symbols', None)}", file=sys.stderr) + except Exception: + pass + + if not args.no_encoder_dump: + import soundfile as sf + wav, sr = sf.read(str(args.wav), dtype="float32", always_2d=False) + if wav.ndim == 2: + wav = wav.mean(axis=1) + if sr != 16000: + import librosa + wav = librosa.resample(wav, orig_sr=sr, target_sr=16000).astype(np.float32) + sr = 16000 + wav_t = torch.from_numpy(wav).unsqueeze(0).to(args.device) + length_t = torch.tensor([len(wav)], dtype=torch.long, device=args.device) + with torch.inference_mode(): + mel, mel_len = model.preprocessor(input_signal=wav_t, length=length_t) + np.save(args.out / "mel.npy", mel[0].detach().cpu().numpy().astype(np.float32)) + enc_out, enc_len = model.encoder(audio_signal=mel, length=mel_len) + enc_np = enc_out[0].permute(1, 0).detach().cpu().numpy().astype(np.float32) + np.save(args.out / "encoder_out.npy", enc_np) + print(f"[rnnt-ref] encoder_out: {enc_np.shape} (T_enc, d_model)", file=sys.stderr) + + print(f"[rnnt-ref] transcribing {args.wav} (NeMo greedy RNN-T)...", file=sys.stderr) + hyps = model.transcribe([str(args.wav)], batch_size=1) + if isinstance(hyps, tuple): + hyps = hyps[0] + h0 = hyps[0] if isinstance(hyps, list) else hyps + + text = h0.text if hasattr(h0, "text") else str(h0) + (args.out / "transcript.txt").write_text(text + "\n", encoding="utf-8") + print(f"[rnnt-ref] transcript: {text!r}", file=sys.stderr) + + token_ids = None + if hasattr(h0, "y_sequence"): + ts = h0.y_sequence + token_ids = (ts.detach().cpu().numpy() if hasattr(ts, "detach") + else np.asarray(ts)).astype(np.int32) + if token_ids is not None: + np.save(args.out / "token_ids.npy", token_ids) + print(f"[rnnt-ref] token_ids: {token_ids.shape} -> token_ids.npy " + f"(first 24: {token_ids[:24].tolist()})", file=sys.stderr) + else: + print("[rnnt-ref] WARN: hypothesis has no y_sequence; no token_ids.npy", file=sys.stderr) + + print(f"[rnnt-ref] done -> {args.out}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/parakeet-cpp/src/main.cpp b/parakeet-cpp/src/main.cpp index 48b0538c7ac..d2dc1054f68 100644 --- a/parakeet-cpp/src/main.cpp +++ b/parakeet-cpp/src/main.cpp @@ -38,16 +38,18 @@ void print_usage(const char * argv0) { PARAKEET_LOG_INFO( "usage: %s --model (--wav | --pcm-in ) [options]\n" "\n" - "Single CLI for all four engine families. The GGUF is auto-detected:\n" + "Single CLI for all five engine families. The GGUF is auto-detected:\n" " CTC (parakeet-ctc-0.6b/1.1b) -> transcription\n" " TDT (parakeet-tdt-0.6b-v3, 1.1b) -> multilingual transcription\n" + " RNNT (plain Transducer head, e.g. the RNN-T branch of a hybrid\n" + " EncDecHybridRNNTCTCBPEModel) -> transcription\n" " EOU (parakeet_realtime_eou_120m-v1) -> low-latency streaming ASR with\n" " native end-of-utterance token\n" " Sortformer (diar_sortformer_4spk-v1, v2) -> 4-speaker diarization\n" "Combined ASR + diarization (\"who said what\") via --diarization-model.\n" "\n" "options:\n" - " --model PATH path to a CTC, TDT, EOU, or Sortformer GGUF (required)\n" + " --model PATH path to a CTC, TDT, RNNT, EOU, or Sortformer GGUF (required)\n" " --wav PATH path to a 16 kHz mono wav file\n" " --pcm-in PATH path to a raw PCM file (mono, format selected by --pcm-format)\n" " --pcm-format FMT raw PCM sample format: s16le (default) or f32le\n" @@ -799,6 +801,23 @@ extern "C" int parakeet_cli_main(int argc, char ** argv) { dopts, dres); rc != 0) return rc; ids_out = std::move(dres.token_ids); text_out = std::move(dres.text); + } else if (model.model_type == ParakeetModelType::RNNT) { + static EouRuntimeWeights rt; + static bool rt_ready = false; + if (!rt_ready) { + if (eou_prepare_runtime(model, rt) != 0) return 20; + rt_ready = true; + } + EouDecodeOptions dopts; + dopts.max_symbols_per_step = model.encoder_cfg.eou_max_symbols_per_step; + dopts.disable_special_tokens = true; // plain greedy RNN-T + EouDecodeResult dres; + if (int rc = eou_greedy_decode(model, rt, + enc_out.encoder_out.data(), + enc_out.n_enc_frames, enc_out.d_model, + dopts, dres); rc != 0) return rc; + ids_out = std::move(dres.token_ids); + text_out = std::move(dres.text); } else { ids_out = ctc_greedy_decode( enc_out.logits.data(), enc_out.n_enc_frames, model.vocab_size, model.blank_id); diff --git a/parakeet-cpp/src/parakeet_ctc.cpp b/parakeet-cpp/src/parakeet_ctc.cpp index 2ef400f3027..16a8510e654 100644 --- a/parakeet-cpp/src/parakeet_ctc.cpp +++ b/parakeet-cpp/src/parakeet_ctc.cpp @@ -1016,6 +1016,7 @@ int load_from_gguf(const std::string & gguf_path, const std::string mtype_str = get_str(g, "parakeet.model.type", "ctc"); if (mtype_str == "tdt") out_model.model_type = ParakeetModelType::TDT; + else if (mtype_str == "rnnt") out_model.model_type = ParakeetModelType::RNNT; else if (mtype_str == "eou") out_model.model_type = ParakeetModelType::EOU; else if (mtype_str == "sortformer") out_model.model_type = ParakeetModelType::SORTFORMER; else out_model.model_type = ParakeetModelType::CTC; @@ -1059,6 +1060,25 @@ int load_from_gguf(const std::string & gguf_path, out_model.eob_id = id_eob >= 0 ? gguf_get_val_i32(g, id_eob) : -1; } + if (out_model.model_type == ParakeetModelType::RNNT) { + // Plain RNN-T (e.g. the Transducer head of a hybrid checkpoint). + // Reuse the encoder_cfg.eou_* predictor/joint fields so the shared + // eou_prepare_runtime + eou_decode_window path works unchanged. There + // are no / tokens: eou_id/eob_id stay -1 and the engine runs + // the decoder in disable_special_tokens mode. vocab_size is the BPE + // label count (no blank); blank_id is that count (blank_as_pad), so + // V_plus_1 = vocab_size + 1 in eou_prepare_runtime. + out_model.encoder_cfg.eou_pred_hidden = get_u32(g, "parakeet.rnnt.pred_hidden", 640); + out_model.encoder_cfg.eou_pred_rnn_layers = get_u32(g, "parakeet.rnnt.pred_rnn_layers", 1); + out_model.encoder_cfg.eou_joint_hidden = get_u32(g, "parakeet.rnnt.joint_hidden", 640); + out_model.encoder_cfg.eou_max_symbols_per_step = get_u32(g, "parakeet.rnnt.max_symbols_per_step", 10); + + out_model.vocab_size = get_u32(g, "parakeet.rnnt.vocab_size", 1024); + out_model.blank_id = get_u32(g, "parakeet.rnnt.blank_id", out_model.vocab_size); + out_model.eou_id = -1; + out_model.eob_id = -1; + } + if (out_model.model_type == ParakeetModelType::SORTFORMER) { out_model.encoder_cfg.sortformer_num_spks = get_u32 (g, "parakeet.sortformer.num_spks", 4); out_model.encoder_cfg.sortformer_fc_d_model = get_u32 (g, "parakeet.sortformer.fc_d_model", 512); @@ -1219,6 +1239,26 @@ int load_from_gguf(const std::string & gguf_path, out_model.eou.joint_pred_b = require_tensor(impl->ctx, "eou.joint.pred.bias"); out_model.eou.joint_out_w = require_tensor(impl->ctx, "eou.joint.out.weight"); out_model.eou.joint_out_b = require_tensor(impl->ctx, "eou.joint.out.bias"); + } else if (out_model.model_type == ParakeetModelType::RNNT) { + // Plain RNN-T predictor + joint, stored in the shared EouWeights slot + // (rnnt.* GGUF tensors; same shapes as eou.* minus the special-token + // rows). Consumed by eou_prepare_runtime via model.eou. + out_model.eou.predict_embed = require_tensor(impl->ctx, "rnnt.predict.embed.weight"); + for (int l = 0; l < out_model.encoder_cfg.eou_pred_rnn_layers; ++l) { + const std::string pl = "rnnt.predict.lstm." + std::to_string(l) + "."; + TdtLstmLayer lyr; + lyr.w_ih = require_tensor(impl->ctx, pl + "w_ih"); + lyr.w_hh = require_tensor(impl->ctx, pl + "w_hh"); + lyr.b_ih = require_tensor(impl->ctx, pl + "b_ih"); + lyr.b_hh = require_tensor(impl->ctx, pl + "b_hh"); + out_model.eou.lstm.push_back(lyr); + } + out_model.eou.joint_enc_w = require_tensor(impl->ctx, "rnnt.joint.enc.weight"); + out_model.eou.joint_enc_b = require_tensor(impl->ctx, "rnnt.joint.enc.bias"); + out_model.eou.joint_pred_w = require_tensor(impl->ctx, "rnnt.joint.pred.weight"); + out_model.eou.joint_pred_b = require_tensor(impl->ctx, "rnnt.joint.pred.bias"); + out_model.eou.joint_out_w = require_tensor(impl->ctx, "rnnt.joint.out.weight"); + out_model.eou.joint_out_b = require_tensor(impl->ctx, "rnnt.joint.out.bias"); } else if (out_model.model_type == ParakeetModelType::SORTFORMER) { out_model.sortformer.encoder_proj_w = require_tensor(impl->ctx, "sortformer.encoder_proj.weight"); out_model.sortformer.encoder_proj_b = require_tensor(impl->ctx, "sortformer.encoder_proj.bias"); @@ -1334,6 +1374,7 @@ ggml_backend_sched_t model_sched(const ParakeetCtcModel & m) { void print_model_summary(const ParakeetCtcModel & m) { const char * mt = "ctc"; if (m.model_type == ParakeetModelType::TDT) mt = "tdt"; + else if (m.model_type == ParakeetModelType::RNNT) mt = "rnnt"; else if (m.model_type == ParakeetModelType::EOU) mt = "eou"; else if (m.model_type == ParakeetModelType::SORTFORMER) mt = "sortformer"; PARAKEET_LOG_INFO("parakeet-%s loaded:\n", mt); @@ -1357,6 +1398,13 @@ void print_model_summary(const ParakeetCtcModel & m) { (double) m.mel_cfg.log_zero_guard_value); if (m.model_type == ParakeetModelType::CTC) { PARAKEET_LOG_INFO(" ctc: vocab=%d blank=%d\n", m.vocab_size, m.blank_id); + } else if (m.model_type == ParakeetModelType::RNNT) { + PARAKEET_LOG_INFO(" rnnt: vocab=%d blank=%d pred_hidden=%d pred_layers=%d " + "joint_hidden=%d max_syms=%d\n", + m.vocab_size, m.blank_id, + m.encoder_cfg.eou_pred_hidden, m.encoder_cfg.eou_pred_rnn_layers, + m.encoder_cfg.eou_joint_hidden, + m.encoder_cfg.eou_max_symbols_per_step); } else if (m.model_type == ParakeetModelType::EOU) { PARAKEET_LOG_INFO(" eou: vocab=%d blank=%d eou_id=%d eob_id=%d " "pred_hidden=%d pred_layers=%d joint_hidden=%d " diff --git a/parakeet-cpp/src/parakeet_ctc.h b/parakeet-cpp/src/parakeet_ctc.h index ed8a86faaf7..d1e39f677e5 100644 --- a/parakeet-cpp/src/parakeet_ctc.h +++ b/parakeet-cpp/src/parakeet_ctc.h @@ -201,6 +201,8 @@ struct TdtWeights { enum class ParakeetModelType { CTC, TDT, + RNNT, // Plain RNN-T (Transducer). Shares EOU's predictor/joint + // runtime + greedy decoder, minus the / tokens. EOU, SORTFORMER, }; diff --git a/parakeet-cpp/src/parakeet_engine.cpp b/parakeet-cpp/src/parakeet_engine.cpp index ce2bbfdd7fe..c517e0f33ab 100644 --- a/parakeet-cpp/src/parakeet_engine.cpp +++ b/parakeet-cpp/src/parakeet_engine.cpp @@ -174,7 +174,10 @@ Engine::Engine(const EngineOptions & opts) : pimpl_(std::make_unique()) { } pimpl_->tdt_ready = true; } - if (pimpl_->model.model_type == ParakeetModelType::EOU) { + if (pimpl_->model.model_type == ParakeetModelType::EOU || + pimpl_->model.model_type == ParakeetModelType::RNNT) { + // RNNT shares EOU's predictor/joint runtime (eou_rt); the decoder runs + // in disable_special_tokens mode for plain greedy RNN-T. if (eou_prepare_runtime(pimpl_->model, pimpl_->eou_rt) != 0) { throw std::runtime_error("Engine: eou_prepare_runtime failed"); } @@ -201,6 +204,7 @@ const EngineOptions & Engine::options() const { std::string Engine::model_type() const { switch (pimpl_->model.model_type) { case ParakeetModelType::TDT: return "tdt"; + case ParakeetModelType::RNNT: return "rnnt"; case ParakeetModelType::EOU: return "eou"; case ParakeetModelType::SORTFORMER: return "sortformer"; case ParakeetModelType::CTC: @@ -215,6 +219,7 @@ bool Engine::is_diarization_model() const { bool Engine::is_transcription_model() const { return pimpl_->model.model_type == ParakeetModelType::CTC || pimpl_->model.model_type == ParakeetModelType::TDT || + pimpl_->model.model_type == ParakeetModelType::RNNT || pimpl_->model.model_type == ParakeetModelType::EOU; } @@ -315,6 +320,20 @@ EngineResult Engine::transcribe_samples(const float * samples, int n_samples, in } ids = std::move(dres.token_ids); text = std::move(dres.text); + } else if (pimpl_->model.model_type == ParakeetModelType::RNNT) { + EouDecodeOptions dopts; + dopts.max_symbols_per_step = pimpl_->model.encoder_cfg.eou_max_symbols_per_step; + dopts.disable_special_tokens = true; // plain greedy RNN-T + EouDecodeResult dres; + if (int rc = eou_greedy_decode(pimpl_->model, pimpl_->eou_rt, + enc_out.encoder_out.data(), + enc_out.n_enc_frames, enc_out.d_model, + dopts, dres); rc != 0) { + throw std::runtime_error("parakeet::Engine::transcribe_samples: rnnt greedy decode failed (rc=" + + std::to_string(rc) + ")"); + } + ids = std::move(dres.token_ids); + text = std::move(dres.text); } else { ids = ctc_greedy_decode(enc_out.logits.data(), enc_out.n_enc_frames, pimpl_->model.vocab_size, pimpl_->model.blank_id); @@ -422,12 +441,13 @@ EngineResult Engine::transcribe_samples_stream(const float * samples, const bool is_tdt = (pimpl_->model.model_type == ParakeetModelType::TDT); const bool is_eou = (pimpl_->model.model_type == ParakeetModelType::EOU); + const bool is_rnnt = (pimpl_->model.model_type == ParakeetModelType::RNNT); int32_t prev_token = -1; TdtDecodeState tdt_state; EouDecodeState eou_state; if (is_tdt) tdt_init_state(pimpl_->tdt_rt, (int) pimpl_->model.blank_id, tdt_state); - if (is_eou) eou_init_state(pimpl_->eou_rt, eou_state); + if (is_eou || is_rnnt) eou_init_state(pimpl_->eou_rt, eou_state); int chunk_index = 0; bool first_segment = true; @@ -470,6 +490,22 @@ EngineResult Engine::transcribe_samples_stream(const float * samples, "eou_decode_window failed (rc=" + std::to_string(rc) + ")"); } eou_boundaries_in_chunk = static_cast(win_segments.size()); + } else if (is_rnnt) { + EouDecodeOptions dopts; + dopts.max_symbols_per_step = pimpl_->model.encoder_cfg.eou_max_symbols_per_step; + dopts.disable_special_tokens = true; // plain greedy RNN-T + std::vector win_segments; + int steps = 0; + const float * win_enc = enc_out.encoder_out.data() + + static_cast(start) * enc_out.d_model; + if (int rc = eou_decode_window(pimpl_->model, pimpl_->eou_rt, + win_enc, end - start, enc_out.d_model, + dopts, eou_state, + win_tokens, win_segments, steps); + rc != 0) { + throw std::runtime_error("parakeet::Engine::transcribe_samples_stream: " + "rnnt decode_window failed (rc=" + std::to_string(rc) + ")"); + } } else { ctc_greedy_decode_window(enc_out.logits.data(), start, end, vocab, blank, @@ -881,7 +917,7 @@ struct StreamSession::Impl { bool finalized = false; bool cancelled = false; - // Optional EnergyVad for CTC/TDT when enable_energy_vad and no native VAD exists. + // Optional EnergyVad for CTC/TDT/RNNT when enable_energy_vad and no native VAD exists. std::unique_ptr energy_vad; int64_t total_pcm_seen = 0; @@ -985,6 +1021,23 @@ void StreamSession::Impl::process_window(const float * window_samples, int windo std::to_string(rc) + ")"); } eou_boundaries_in_chunk = static_cast(win_segments.size()); + } else if (engine_impl->model.model_type == ParakeetModelType::RNNT) { + EouDecodeOptions dopts; + dopts.max_symbols_per_step = engine_impl->model.encoder_cfg.eou_max_symbols_per_step; + dopts.disable_special_tokens = true; // plain greedy RNN-T + std::vector win_segments; + int steps = 0; + const int n_frames = std::max(0, center_end_frame - left_drop_frames); + const float * win_enc = enc_out.encoder_out.data() + + static_cast(left_drop_frames) * enc_out.d_model; + if (int rc = eou_decode_window(engine_impl->model, engine_impl->eou_rt, + win_enc, n_frames, enc_out.d_model, + dopts, eou_state, + win_tokens, win_segments, steps); + rc != 0) { + throw std::runtime_error("StreamSession: rnnt decode_window failed (rc=" + + std::to_string(rc) + ")"); + } } else { ctc_greedy_decode_window(enc_out.logits.data(), left_drop_frames, center_end_frame, @@ -1197,7 +1250,9 @@ std::unique_ptr Engine::stream_start(const StreamingOptions & opt if (pimpl_->model.model_type == ParakeetModelType::TDT) { tdt_init_state(pimpl_->tdt_rt, (int) pimpl_->model.blank_id, impl->tdt_state); } - // Optional EnergyVad for CTC/TDT only (EOU uses ``; Sortformer uses SortformerStreamSession). + // Optional EnergyVad for CTC/TDT/RNNT (EOU uses ``; Sortformer uses + // SortformerStreamSession). Plain RNNT has no native end-pointing signal, + // so it takes the same opt-in RMS VAD as CTC/TDT. if (opts.enable_energy_vad && pimpl_->model.model_type != ParakeetModelType::EOU) { impl->energy_vad = std::make_unique( @@ -1206,7 +1261,8 @@ std::unique_ptr Engine::stream_start(const StreamingOptions & opt opts.energy_vad_hangover_ms, opts.energy_vad_threshold_db); } - if (pimpl_->model.model_type == ParakeetModelType::EOU) { + if (pimpl_->model.model_type == ParakeetModelType::EOU || + pimpl_->model.model_type == ParakeetModelType::RNNT) { eou_init_state(pimpl_->eou_rt, impl->eou_state); } diff --git a/parakeet-cpp/src/parakeet_eou.cpp b/parakeet-cpp/src/parakeet_eou.cpp index 84bcd1e0d4c..708b139c016 100644 --- a/parakeet-cpp/src/parakeet_eou.cpp +++ b/parakeet-cpp/src/parakeet_eou.cpp @@ -170,7 +170,8 @@ std::string trim_spaces(const std::string & s) { } int eou_prepare_runtime(const ParakeetCtcModel & model, EouRuntimeWeights & W) { - if (model.model_type != ParakeetModelType::EOU) { + const bool is_rnnt = model.model_type == ParakeetModelType::RNNT; + if (model.model_type != ParakeetModelType::EOU && !is_rnnt) { return 1; } W.H_pred = model.encoder_cfg.eou_pred_hidden; @@ -179,8 +180,11 @@ int eou_prepare_runtime(const ParakeetCtcModel & model, EouRuntimeWeights & W) { W.L = model.encoder_cfg.eou_pred_rnn_layers; W.V_plus_1 = (int) model.vocab_size + 1; W.blank_id = (int) model.blank_id; - W.eou_id = model.eou_id >= 0 ? model.eou_id : (int) model.vocab_size - 2; - W.eob_id = model.eob_id >= 0 ? model.eob_id : (int) model.vocab_size - 1; + // Plain RNN-T has no /: keep the sentinels at -1 (never matched) + // rather than EOU's vocab_size-2 / -1 fallback, which would alias real BPE + // tokens. The engine also runs the decoder in disable_special_tokens mode. + W.eou_id = is_rnnt ? -1 : (model.eou_id >= 0 ? model.eou_id : (int) model.vocab_size - 2); + W.eob_id = is_rnnt ? -1 : (model.eob_id >= 0 ? model.eob_id : (int) model.vocab_size - 1); dequantize_to_f32(model.eou.predict_embed, W.embed); @@ -254,6 +258,7 @@ int eou_decode_window(const ParakeetCtcModel & model, const int eou = W.eou_id; const int eob = W.eob_id; const int max_syms = std::max(1, opts.max_symbols_per_step); + const bool plain = opts.disable_special_tokens; // plain RNN-T: blank-only break std::vector scratch_lstm; std::vector scratch_lstm_layer_input; @@ -280,7 +285,7 @@ int eou_decode_window(const ParakeetCtcModel & model, } // : training-time block boundary; treat as a no-op skip. - if (best == eob) { + if (!plain && best == eob) { break; } @@ -289,7 +294,7 @@ int eou_decode_window(const ParakeetCtcModel & model, // NeMo `eouDecodeChunk` reference: do NOT feed `` // back into the predictor; reset h/c to zero and lastToken // to blank. - if (best == eou) { + if (!plain && best == eou) { if (state.has_emitted_token_since_last_eou) { EouSegmentBoundary boundary; boundary.token_index = (int) out_tokens.size(); @@ -312,7 +317,7 @@ int eou_decode_window(const ParakeetCtcModel & model, // Skip any other special token defensively (e.g. ); // any vocab piece wrapped in `<...>` is treated as special. - if (best >= 0 && (size_t) best < n_vocab) { + if (!plain && best >= 0 && (size_t) best < n_vocab) { const std::string & piece = model.vocab.pieces[best]; if (!piece.empty() && piece.front() == '<' && piece.back() == '>') { break; diff --git a/parakeet-cpp/src/parakeet_eou.h b/parakeet-cpp/src/parakeet_eou.h index c21fbdec008..5a2887000fd 100644 --- a/parakeet-cpp/src/parakeet_eou.h +++ b/parakeet-cpp/src/parakeet_eou.h @@ -48,6 +48,9 @@ struct EouRuntimeWeights { struct EouDecodeOptions { int max_symbols_per_step = 5; + // Plain RNN-T mode: disable all //<...> special-token handling so + // the greedy loop breaks only on the transducer blank (matches NeMo greedy). + bool disable_special_tokens = false; }; struct EouDecodeState { diff --git a/parakeet-cpp/test/test_encoder_capture_parity.cpp b/parakeet-cpp/test/test_encoder_capture_parity.cpp index 70a4975a8cb..a1f39ef55f7 100644 --- a/parakeet-cpp/test/test_encoder_capture_parity.cpp +++ b/parakeet-cpp/test/test_encoder_capture_parity.cpp @@ -79,16 +79,17 @@ int main(int argc, char ** argv) { } // Capture-parity gate works on any model type. CTC GGUFs populate - // both `encoder_out` and `logits`; TDT/EOU/Sortformer GGUFs only + // both `encoder_out` and `logits`; TDT/RNNT/EOU/Sortformer GGUFs only // populate `encoder_out` (their decoders consume `encoder_out` and // produce their own logits separately). For those, `logits` is // empty in BOTH calls, so the byte-equal check trivially holds — // we keep it in the assertion path so any future change that - // accidentally starts populating logits on a TDT/EOU/Sortformer + // accidentally starts populating logits on a TDT/RNNT/EOU/Sortformer // path will be caught. const char * mt_name = model.model_type == ParakeetModelType::CTC ? "ctc" : model.model_type == ParakeetModelType::TDT ? "tdt" + : model.model_type == ParakeetModelType::RNNT ? "rnnt" : model.model_type == ParakeetModelType::EOU ? "eou" : model.model_type == ParakeetModelType::SORTFORMER ? "sortformer" : "unknown"; diff --git a/parakeet-cpp/test/test_rnnt_decoder_parity.cpp b/parakeet-cpp/test/test_rnnt_decoder_parity.cpp new file mode 100644 index 00000000000..946094ad725 --- /dev/null +++ b/parakeet-cpp/test/test_rnnt_decoder_parity.cpp @@ -0,0 +1,268 @@ +// Plain RNN-T decoder parity vs reference token IDs (NeMo dump or cross-backend). +// +// Greedy decoding is deterministic; this compares integer token IDs only. +// +// Usage: +// test-rnnt-decoder-parity [] +// +// Pass containing token_ids.npy from +// `scripts/dump-rnnt-reference.py --wav ` +// to assert bit-identical greedy decoding against NeMo. +// +// Exit 0 on success; non-zero on failure or invalid arguments. + +#include "parakeet_ctc.h" +#include "parakeet_eou.h" +#include "mel_preprocess.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +int load_npy_i32(const std::string & path, + std::vector & out_data) { + std::ifstream f(path, std::ios::binary); + if (!f) return 1; + + char magic[6]; + f.read(magic, 6); + if (std::memcmp(magic, "\x93NUMPY", 6) != 0) return 2; + + uint8_t major = 0, minor = 0; + f.read(reinterpret_cast(&major), 1); + f.read(reinterpret_cast(&minor), 1); + + uint32_t header_len = 0; + if (major == 1) { + uint16_t hl = 0; + f.read(reinterpret_cast(&hl), 2); + header_len = hl; + } else { + f.read(reinterpret_cast(&header_len), 4); + } + + std::string header(header_len, '\0'); + f.read(header.data(), header_len); + + // Little-endian int32, C order only -- a reference saved without + // .astype(np.int32) (numpy defaults to int64) would otherwise parse + // "successfully" and compare as interleaved garbage. + if (header.find("'descr': '= shape_str.size()) break; + size_t end = pos; + while (end < shape_str.size() && + std::isdigit(static_cast(shape_str[end]))) ++end; + if (end > pos) { + total *= static_cast(std::stoll(shape_str.substr(pos, end - pos))); + any = true; + pos = end; + } else { + break; + } + } + if (!any) return 4; + + out_data.resize(total); + f.read(reinterpret_cast(out_data.data()), + total * sizeof(int32_t)); + return f ? 0 : 5; +} + +// Run the full pipeline (mel -> encoder -> plain RNN-T greedy) on the given +// GGUF using the requested n_gpu_layers, returning token IDs + transcript. +int transcribe_rnnt(const std::string & gguf_path, + const std::string & wav_path, + int n_gpu_layers, + std::vector & out_tokens, + std::string & out_text) { + using namespace parakeet; + + ParakeetCtcModel model; + if (int rc = load_from_gguf(gguf_path, model, /*n_threads=*/0, + n_gpu_layers, /*verbose=*/false); rc != 0) { + std::fprintf(stderr, " load_from_gguf failed rc=%d\n", rc); + return 100 + rc; + } + if (model.model_type != ParakeetModelType::RNNT) { + std::fprintf(stderr, " error: expected RNNT model in %s\n", + gguf_path.c_str()); + return 110; + } + + std::vector samples; + int sr = 0; + if (int rc = load_wav_mono_f32(wav_path, samples, sr); rc != 0) { + std::fprintf(stderr, " load_wav failed rc=%d\n", rc); + return 120 + rc; + } + if (sr != model.mel_cfg.sample_rate) { + std::fprintf(stderr, + " error: wav is %d Hz but the model expects %d Hz " + "(this harness does not resample)\n", + sr, model.mel_cfg.sample_rate); + return 125; + } + + std::vector mel; + int n_frames = 0; + if (int rc = compute_log_mel(samples.data(), (int) samples.size(), + model.mel_cfg, mel, n_frames); rc != 0) { + std::fprintf(stderr, " compute_log_mel failed rc=%d\n", rc); + return 130 + rc; + } + + EncoderOutputs enc_out; + if (int rc = run_encoder(model, mel.data(), n_frames, + model.mel_cfg.n_mels, enc_out); rc != 0) { + std::fprintf(stderr, " run_encoder failed rc=%d\n", rc); + return 140 + rc; + } + + EouRuntimeWeights rt; + if (int rc = eou_prepare_runtime(model, rt); rc != 0) { + std::fprintf(stderr, " eou_prepare_runtime failed rc=%d\n", rc); + return 150 + rc; + } + + EouDecodeOptions dopts; + dopts.max_symbols_per_step = model.encoder_cfg.eou_max_symbols_per_step; + dopts.disable_special_tokens = true; // plain greedy RNN-T + EouDecodeResult dres; + if (int rc = eou_greedy_decode(model, rt, + enc_out.encoder_out.data(), + enc_out.n_enc_frames, enc_out.d_model, + dopts, dres); rc != 0) { + std::fprintf(stderr, " eou_greedy_decode failed rc=%d\n", rc); + return 160 + rc; + } + + out_tokens = std::move(dres.token_ids); + out_text = std::move(dres.text); + return 0; +} + +void print_first_diff(const std::vector & a, + const std::vector & b) { + const size_t n = std::min(a.size(), b.size()); + for (size_t i = 0; i < n; ++i) { + if (a[i] != b[i]) { + std::fprintf(stderr, " first diff at index %zu: a=%d b=%d\n", + i, a[i], b[i]); + return; + } + } + std::fprintf(stderr, " no per-index diff in shared prefix; lengths %zu vs %zu\n", + a.size(), b.size()); +} + +} // namespace + +int main(int argc, char ** argv) { + if (argc < 3) { + std::fprintf(stderr, + "usage: %s []\n" + "\n" + "Validates the C++ plain RNN-T greedy decoder.\n" + " Pass containing token_ids.npy from\n" + " `scripts/dump-rnnt-reference.py --wav `\n" + " to compare against the NeMo greedy reference.\n" + "\n" + "Always cross-checks n_gpu_layers=0 against n_gpu_layers=1 (encoder\n" + "on the compiled-in GPU backend when one is available; the RNN-T\n" + "greedy decode itself is scalar CPU in both configurations).\n", + argv[0]); + return 2; + } + + const std::string gguf_path = argv[1]; + const std::string wav_path = argv[2]; + const std::string ref_dir = (argc >= 4) ? argv[3] : ""; + + std::fprintf(stderr, "[rnnt-decode-parity] running CPU fallback (n_gpu_layers=0)...\n"); + std::vector ids_cpu; + std::string text_cpu; + if (int rc = transcribe_rnnt(gguf_path, wav_path, 0, ids_cpu, text_cpu); rc != 0) { + return rc; + } + std::fprintf(stderr, "[rnnt-decode-parity] CPU: tokens=%zu text=%.80s%s\n", + ids_cpu.size(), text_cpu.c_str(), + text_cpu.size() > 80 ? "..." : ""); + + std::fprintf(stderr, "[rnnt-decode-parity] running GPU-offloaded encoder (n_gpu_layers=1)...\n"); + std::vector ids_gpu; + std::string text_gpu; + if (int rc = transcribe_rnnt(gguf_path, wav_path, 1, ids_gpu, text_gpu); rc != 0) { + return rc; + } + std::fprintf(stderr, "[rnnt-decode-parity] GPU: tokens=%zu text=%.80s%s\n", + ids_gpu.size(), text_gpu.c_str(), + text_gpu.size() > 80 ? "..." : ""); + + bool ok = true; + if (ids_cpu.size() != ids_gpu.size() || + !std::equal(ids_cpu.begin(), ids_cpu.end(), ids_gpu.begin())) { + std::fprintf(stderr, + "[rnnt-decode-parity] FAIL: CPU vs GPU-encoder token IDs differ\n"); + print_first_diff(ids_cpu, ids_gpu); + ok = false; + } else { + std::fprintf(stderr, + "[rnnt-decode-parity] PASS: CPU vs GPU-encoder token IDs match (%zu tokens)\n", + ids_cpu.size()); + } + + if (!ref_dir.empty()) { + std::vector ids_ref; + const std::string p = ref_dir + "/token_ids.npy"; + if (int rc = load_npy_i32(p, ids_ref); rc != 0) { + // A ref-dir was explicitly requested: an unreadable reference is + // a failure, not a skip -- otherwise a typo'd path silently + // un-asserts the only NeMo comparison this harness exists for. + std::fprintf(stderr, + "[rnnt-decode-parity] FAIL: could not load %s (rc=%d)\n", + p.c_str(), rc); + ok = false; + } else { + if (ids_ref.size() != ids_gpu.size() || + !std::equal(ids_ref.begin(), ids_ref.end(), ids_gpu.begin())) { + std::fprintf(stderr, + "[rnnt-decode-parity] FAIL: NeMo reference (%zu tokens) vs C++ (%zu) mismatch\n", + ids_ref.size(), ids_gpu.size()); + print_first_diff(ids_ref, ids_gpu); + ok = false; + } else { + std::fprintf(stderr, + "[rnnt-decode-parity] PASS: NeMo reference matches (%zu tokens)\n", + ids_ref.size()); + } + } + } + + if (ok) { + std::fprintf(stderr, "[rnnt-decode-parity] all checks passed\n"); + return 0; + } + std::fprintf(stderr, "[rnnt-decode-parity] one or more checks failed\n"); + return 1; +}