diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 04467e0..c844881 100644 --- a/.github/workflows/ci.yml +++ b/.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 \ diff --git a/AGENTS.md b/AGENTS.md index f512a53..fa3334b 100644 --- a/AGENTS.md +++ b/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 diff --git a/README.md b/README.md index 74a9503..a71c2ae 100644 --- a/README.md +++ b/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`. --- diff --git a/docs/tdt-nbest.md b/docs/tdt-nbest.md new file mode 100644 index 0000000..6800da1 --- /dev/null +++ b/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. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 69a2b2d..0db3f1a 100644 --- a/examples/cli/main.cpp +++ b/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,13 +222,20 @@ 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 --input " "[--decoder ctc|tdt] [--lang ] [--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 @@ -233,6 +243,11 @@ static int cmd_transcribe(int argc, char** argv) { 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 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 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 \n" " parakeet-cli transcribe --model --input " "[--decoder ctc|tdt] [--lang ] [--stream] [--timestamps] " - "[--threads N] [--json]\n" + "[--threads N] [--json] " + "[--beam-size N [--nbest N] [--no-score-norm]]\n" " parakeet-cli quantize " "\n" " parakeet-cli bench --model --manifest " diff --git a/include/parakeet_capi.h b/include/parakeet_capi.h index c199a70..e84e485 100644 --- a/include/parakeet_capi.h +++ b/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 diff --git a/scripts/gen_nemo_baseline.py b/scripts/gen_nemo_baseline.py index e5b717e..9b2e284 100644 --- a/scripts/gen_nemo_baseline.py +++ b/scripts/gen_nemo_baseline.py @@ -83,6 +83,14 @@ ``m.transcribe([audio])``. May be length-0 for a silent / tone clip. +Optional TDT N-best ground truth (``--tdt-beam-size N``): + +* ``tdt_nbest_offsets`` ``[N+1]`` int32 ragged-hypothesis offsets +* ``tdt_nbest_token_ids`` ``[sum L]`` int32 flattened token ids +* ``tdt_nbest_token_end_frames`` ``[sum L]`` int32 NeMo token end frames +* ``tdt_nbest_scores`` ``[N]`` f32 raw hypothesis scores +* ``baseline.tdt_nbest_texts_json`` str decoded texts as JSON + String KVs: * ``baseline.detok_text`` detokenized text for the ``detok_ids`` fixture @@ -441,6 +449,62 @@ def _timestamps_decoding_cfg(m): return cfg +def _dump_tdt_nbest(m, audio, beam_size): + """Run NeMo's default TDT beam search and return flat parity arrays.""" + import copy + from omegaconf import open_dict + + m.change_decoding_strategy(decoder_type="rnnt") + cfg = copy.deepcopy(m.cfg.decoding) + with open_dict(cfg): + cfg.strategy = "beam" + cfg.compute_timestamps = False + cfg.preserve_alignments = False + cfg.beam.beam_size = beam_size + cfg.beam.score_norm = True + cfg.beam.return_best_hypothesis = False + m.change_decoding_strategy(cfg, decoder_type="rnnt") + + with torch.no_grad(): + out = m.transcribe( + [audio], batch_size=1, return_hypotheses=True + ) + samples = out[0] if isinstance(out, tuple) else out + hypotheses = samples[0] + if not isinstance(hypotheses, list): + hypotheses = [hypotheses] + hypotheses = hypotheses[:beam_size] + + offsets = [0] + token_ids = [] + token_end_frames = [] + scores = [] + texts = [] + for hyp in hypotheses: + ids = hyp.y_sequence + ids = ids.cpu().tolist() if hasattr(ids, "cpu") else list(ids) + ends = hyp.timestamp + ends = ends.cpu().tolist() if hasattr(ends, "cpu") else list(ends) + if len(ids) != len(ends): + raise RuntimeError( + "NeMo TDT N-best token/timestamp length mismatch: " + f"{len(ids)} != {len(ends)}" + ) + token_ids.extend(int(x) for x in ids) + token_end_frames.extend(int(x) for x in ends) + offsets.append(len(token_ids)) + scores.append(float(hyp.score)) + texts.append(hyp.text if hasattr(hyp, "text") else str(hyp)) + + return { + "offsets": np.asarray(offsets, dtype=np.int32), + "token_ids": np.asarray(token_ids, dtype=np.int32), + "token_end_frames": np.asarray(token_end_frames, dtype=np.int32), + "scores": np.asarray(scores, dtype=np.float32), + "texts": texts, + } + + def _dump_head_timestamps(m, audio, decoder_type, blank_id): """Run one decoding head with timestamps+confidence and extract the per-token ``{id, frame, conf}`` arrays plus the per-word ``{w, start, end, conf}`` list. @@ -710,7 +774,16 @@ def main(): default=None, help="target_lang for prompt models (default: model default / auto)", ) + ap.add_argument( + "--tdt-beam-size", + type=int, + default=0, + help="also dump NeMo default TDT N-best parity tensors using this " + "beam size (0 disables; intended for beam decoder tests)", + ) args = ap.parse_args() + if args.tdt_beam_size < 0: + ap.error("--tdt-beam-size must be >= 0") is_local = pathlib.Path(args.model).exists() try: @@ -1037,6 +1110,11 @@ def submodule(path): y_seq = y_seq.cpu().tolist() tdt_token_ids = np.array(list(y_seq), dtype=np.int32) tdt_text = tdt_first.text if hasattr(tdt_first, "text") else str(tdt_first) + tdt_nbest = ( + _dump_tdt_nbest(m, args.audio, args.tdt_beam_size) + if args.tdt_beam_size > 0 + else None + ) # Tokenizer detok fixture: a small hand-picked set of regular BPE ids # (avoid id=0 and blank_id=1024 which is beyond vocab). @@ -1071,6 +1149,28 @@ def submodule(path): w.add_uint32("baseline.tdt_token_count", int(tdt_token_ids.shape[0])) if tdt_token_ids.shape[0] > 0: w.add_tensor("tdt_token_ids", np.ascontiguousarray(tdt_token_ids)) + if tdt_nbest is not None: + w.add_uint32("baseline.tdt_nbest_beam_size", args.tdt_beam_size) + w.add_uint32( + "baseline.tdt_nbest_count", int(tdt_nbest["scores"].shape[0]) + ) + w.add_tensor( + "tdt_nbest_offsets", np.ascontiguousarray(tdt_nbest["offsets"]) + ) + w.add_tensor( + "tdt_nbest_token_ids", np.ascontiguousarray(tdt_nbest["token_ids"]) + ) + w.add_tensor( + "tdt_nbest_token_end_frames", + np.ascontiguousarray(tdt_nbest["token_end_frames"]), + ) + w.add_tensor( + "tdt_nbest_scores", np.ascontiguousarray(tdt_nbest["scores"]) + ) + w.add_string( + "baseline.tdt_nbest_texts_json", + json.dumps(tdt_nbest["texts"], ensure_ascii=False), + ) w.add_string("baseline.detok_text", detok_text) w.add_string("baseline.ctc_text", ctc_text) w.add_string("baseline.tdt_text", tdt_text) @@ -1086,11 +1186,25 @@ def submodule(path): shapes["joint_enc_frames"] = tuple(joint_enc_frames_arr.shape) if tdt_token_ids.shape[0] > 0: shapes["tdt_token_ids"] = tuple(tdt_token_ids.shape) + if tdt_nbest is not None: + for name in ( + "offsets", + "token_ids", + "token_end_frames", + "scores", + ): + shapes[f"tdt_nbest_{name}"] = tuple(tdt_nbest[name].shape) print("baseline tensors:", shapes) print(f"baseline.detok_text: {repr(detok_text)}") print(f"baseline.ctc_text: {repr(ctc_text)}") print(f"baseline.tdt_text: {repr(tdt_text)}") print(f"tdt_token_ids ({tdt_token_ids.shape[0]}): {tdt_token_ids.tolist()}") + if tdt_nbest is not None: + print( + f"tdt_nbest beam={args.tdt_beam_size}: " + f"scores={tdt_nbest['scores'].tolist()} " + f"texts={tdt_nbest['texts']!r}" + ) print( f"transducer: pred_input_ids={pred_input_ids.tolist()} add_sos={add_sos} " f"U+1={shapes['pred_out'][0]} sos_embed_zero={sos_embed_is_zero} " diff --git a/src/model.cpp b/src/model.cpp index af38fb5..e3ae0c1 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -75,6 +75,67 @@ static void maybe_apply_prompt(const ModelLoader& loader, std::vector& en enc_out.swap(projected); } +struct EncodedAudio { + std::vector channels_first; + int d_model = 0; + int frames = 0; +}; + +static void validate_nbest_request(const ParakeetConfig& cfg, + int beam_size, int nbest) { + if (cfg.tdt_durations.empty()) + throw std::runtime_error( + "parakeet: N-best decoding requires a TDT duration table"); + if (beam_size < 1 || nbest < 1 || nbest > beam_size) + throw std::invalid_argument( + "parakeet: require beam_size >= nbest >= 1"); +} + +// Shared single-clip frontend + encoder path. Keeping this in one helper makes +// the opt-in N-best decoder consume exactly the same encoder output as greedy +// and timestamped transcription. +static EncodedAudio encode_16k(const ModelLoader& loader, + const std::vector& pcm16k, + int prompt_index) { + const ParakeetConfig& cfg = loader.config(); + + std::vector feats; + int n_mels = 0, T = 0; + if (std::string(pk::global_backend().device_name()) != "cpu") { + GpuMel gmel(loader); + gmel.compute(pcm16k, feats, n_mels, T); + } else { + MelFrontend mel(loader); + mel.compute(pcm16k, feats, n_mels, T); + } + + Encoder encoder(loader); + EncodedAudio encoded; + const int sub_tile = subsampling_tile_for(cfg, loader, T); + if (sub_tile > 0) { + MelBatch mb1; + mb1.B = 1; + mb1.n_mels = n_mels; + mb1.T_max = T; + mb1.valid_T = {T}; + mb1.data = feats; + std::vector> outputs; + std::vector valid_frames; + int padded_frames = 0; + encoder.forward_batch_tiled( + mb1, outputs, encoded.d_model, padded_frames, valid_frames, sub_tile); + encoded.channels_first = std::move(outputs[0]); + encoded.frames = valid_frames[0]; + } else { + encoder.forward(feats, n_mels, T, encoded.channels_first, + encoded.d_model, encoded.frames); + } + + maybe_apply_prompt(loader, encoded.channels_first, + encoded.d_model, encoded.frames, prompt_index); + return encoded; +} + // Decode one item's encoder output (row-major [d_model, Tout], channels-first) // into a transcript. Mirrors the tail of transcribe_16k exactly. static std::string decode_enc_out(const ModelLoader& loader, @@ -112,52 +173,14 @@ std::string Model::transcribe_16k(const std::vector& pcm16k, const std::string& target_lang) const { const ParakeetConfig& cfg = loader_.config(); const int prompt_index = resolve_prompt_index(target_lang); - - // 1. Log-mel front end -> feats [n_mels, T]. On a non-CPU backend run the - // heavy STFT/power/filterbank/log on the backend (GPU) via GpuMel; on CPU - // keep the byte-identical FFT-based MelFrontend. - std::vector feats; - int n_mels = 0, T = 0; - if (std::string(pk::global_backend().device_name()) != "cpu") { - GpuMel gmel(loader_); - gmel.compute(pcm16k, feats, n_mels, T); - } else { - MelFrontend mel(loader_); - mel.compute(pcm16k, feats, n_mels, T); - } - - // 2. FastConformer encoder -> enc_out [d_model, Tout] (channels-first). - Encoder encoder(loader_); - std::vector enc_out; - int d_model = 0, Tout = 0; - // Long audio: the single-clip forward() would overflow ggml's 2^31 subsampling - // limit. Delegate to the tiled batched path with a 1-item batch (faithful, - // reuses forward_batch_tiled). Short audio keeps the fused single-clip path. - const int sub_tile = subsampling_tile_for(cfg, loader_, T); - if (sub_tile > 0) { - MelBatch mb1; - mb1.B = 1; mb1.n_mels = n_mels; mb1.T_max = T; mb1.valid_T = { T }; - mb1.data = feats; // feats is [n_mels,T] = the B=1 batch buffer - std::vector> eo; std::vector vT; - int dm1 = 0, To1 = 0; - encoder.forward_batch_tiled(mb1, eo, dm1, To1, vT, sub_tile); - enc_out = std::move(eo[0]); // channels-first [d_model, vT[0]] - d_model = dm1; - Tout = vT[0]; - } else { - encoder.forward(feats, n_mels, T, enc_out, d_model, Tout); - } - - // 2b. Prompt conditioning (multilingual nemotron): project the encoder - // output with the selected language one-hot before decoding. No-op for - // other models (prompt.present == false). - maybe_apply_prompt(loader_, enc_out, d_model, Tout, prompt_index); + EncodedAudio encoded = encode_16k(loader_, pcm16k, prompt_index); // Decide which head to use. const bool use_tdt = (decoder == Decoder::kTDT) || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); - return decode_enc_out(loader_, enc_out, d_model, Tout, use_tdt); + return decode_enc_out(loader_, encoded.channels_first, + encoded.d_model, encoded.frames, use_tdt); } // Max mel frames per encoder pass before the first subsampling conv output @@ -365,50 +388,14 @@ Transcription Model::transcribe_16k_with_timestamps( // hop_length / sample_rate). const float frame_sec = (float)cfg.hop_length * (float)cfg.subsampling_factor / (float)cfg.sample_rate; - - // 1. Log-mel front end -> feats [n_mels, T]. On a non-CPU backend run the - // heavy STFT/power/filterbank/log on the backend (GPU) via GpuMel; on CPU - // keep the byte-identical FFT-based MelFrontend. - std::vector feats; - int n_mels = 0, T = 0; - if (std::string(pk::global_backend().device_name()) != "cpu") { - GpuMel gmel(loader_); - gmel.compute(pcm16k, feats, n_mels, T); - } else { - MelFrontend mel(loader_); - mel.compute(pcm16k, feats, n_mels, T); - } - - // 2. FastConformer encoder -> enc_out [d_model, Tout] (channels-first). - Encoder encoder(loader_); - std::vector enc_out; - int d_model = 0, Tout = 0; - // Long audio: the single-clip forward() would overflow ggml's 2^31 subsampling - // limit. Delegate to the tiled batched path with a 1-item batch (faithful, - // reuses forward_batch_tiled). Short audio keeps the fused single-clip path. - const int sub_tile = subsampling_tile_for(cfg, loader_, T); - if (sub_tile > 0) { - MelBatch mb1; - mb1.B = 1; mb1.n_mels = n_mels; mb1.T_max = T; mb1.valid_T = { T }; - mb1.data = feats; // feats is [n_mels,T] = the B=1 batch buffer - std::vector> eo; std::vector vT; - int dm1 = 0, To1 = 0; - encoder.forward_batch_tiled(mb1, eo, dm1, To1, vT, sub_tile); - enc_out = std::move(eo[0]); // channels-first [d_model, vT[0]] - d_model = dm1; - Tout = vT[0]; - } else { - encoder.forward(feats, n_mels, T, enc_out, d_model, Tout); - } - - // 2b. Prompt conditioning (nemotron): project before decode. No-op otherwise. - maybe_apply_prompt(loader_, enc_out, d_model, Tout, prompt_index); + EncodedAudio encoded = encode_16k(loader_, pcm16k, prompt_index); const bool use_tdt = (decoder == Decoder::kTDT) || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); Transcription result = decode_enc_out_with_timestamps( - loader_, enc_out, d_model, Tout, use_tdt, frame_sec); + loader_, encoded.channels_first, encoded.d_model, encoded.frames, + use_tdt, frame_sec); return result; } @@ -487,6 +474,75 @@ std::vector Model::transcribe_pcm_batch_with_timestamps( return transcribe_16k_batch_with_timestamps(r, decoder, target_lang); } +std::vector Model::transcribe_16k_nbest( + const std::vector& pcm16k, int beam_size, int nbest, + bool score_norm, const std::string& target_lang) const { + const ParakeetConfig& cfg = loader_.config(); + + const int prompt_index = resolve_prompt_index(target_lang); + EncodedAudio encoded = encode_16k(loader_, pcm16k, prompt_index); + + std::vector enc_row( + (size_t)encoded.frames * encoded.d_model); + for (int t = 0; t < encoded.frames; ++t) + for (int c = 0; c < encoded.d_model; ++c) + enc_row[(size_t)t * encoded.d_model + c] = + encoded.channels_first[(size_t)c * encoded.frames + t]; + + PredictionNet pred(loader_); + Joint joint(loader_); + std::vector beam = tdt_beam_search( + pred, joint, enc_row, encoded.frames, encoded.d_model, + cfg.tdt_durations, (int)cfg.blank_id, + beam_size, nbest, score_norm); + + std::vector result; + result.reserve(beam.size()); + for (TdtBeamHypothesis& hyp : beam) { + std::vector ids; + ids.reserve(hyp.tokens.size()); + for (const TdtBeamToken& token : hyp.tokens) + ids.push_back(token.id); + + NBestTranscription item; + item.text = detokenize( + loader_.tokenizer_pieces(), + strip_special_tokens(loader_.tokenizer_pieces(), ids)); + item.tokens = std::move(hyp.tokens); + item.score = hyp.score; + item.normalized_score = hyp.normalized_score; + result.push_back(std::move(item)); + } + return result; +} + +std::vector Model::transcribe_pcm_nbest( + const std::vector& pcm, int sample_rate, + int beam_size, int nbest, bool score_norm, + const std::string& target_lang) const { + validate_nbest_request(loader_.config(), beam_size, nbest); + if (sample_rate <= 0) + throw std::runtime_error("parakeet: invalid sample_rate"); + if (sample_rate == 16000) + return transcribe_16k_nbest( + pcm, beam_size, nbest, score_norm, target_lang); + return transcribe_16k_nbest( + resample_linear(pcm, sample_rate, 16000), + beam_size, nbest, score_norm, target_lang); +} + +std::vector Model::transcribe_path_nbest( + const std::string& wav_path, int beam_size, int nbest, + bool score_norm, const std::string& target_lang) const { + validate_nbest_request(loader_.config(), beam_size, nbest); + Audio audio; + if (!load_audio_16k_mono(wav_path, audio)) + throw std::runtime_error( + "parakeet: failed to load audio: " + wav_path); + return transcribe_16k_nbest( + audio.samples, beam_size, nbest, score_norm, target_lang); +} + std::string Model::transcribe_pcm(const std::vector& pcm, int sample_rate, Decoder decoder, const std::string& target_lang) const { if (sample_rate <= 0) { diff --git a/src/model.hpp b/src/model.hpp index 30adeaa..4ff6fa4 100644 --- a/src/model.hpp +++ b/src/model.hpp @@ -1,6 +1,7 @@ #pragma once #include "parakeet.h" // pk::Decoder #include "model_loader.hpp" +#include "tdt.hpp" // pk::TdtBeamToken #include "transcription.hpp" // pk::Transcription #include @@ -9,6 +10,14 @@ namespace pk { +// One text hypothesis returned by the opt-in offline TDT N-best path. +struct NBestTranscription { + std::string text; + std::vector tokens; + float score = 0.0f; + float normalized_score = 0.0f; +}; + // Load-once transcription context. // // Loads a GGUF model ONCE (owns the ModelLoader) and reuses it across many @@ -76,6 +85,19 @@ class Model { Decoder decoder = Decoder::kDefault, const std::string& target_lang = "") const; + // Offline TDT beam search. These methods are intentionally separate from + // the greedy Decoder selector: they require a TDT duration table and return + // up to `nbest` ranked hypotheses with token emission frames/durations. + std::vector transcribe_pcm_nbest( + const std::vector& pcm, int sample_rate, + int beam_size, int nbest, bool score_norm = true, + const std::string& target_lang = "") const; + + std::vector transcribe_path_nbest( + const std::string& wav_path, + int beam_size, int nbest, bool score_norm = true, + const std::string& target_lang = "") const; + // Batched timestamped transcription. Each clip is resampled to 16 kHz if // needed, all run through the batched encoder; decode + timestamp extraction // are per item. Returns one Transcription per input, in order. @@ -115,6 +137,10 @@ class Model { const std::vector& pcm16k, Decoder decoder, const std::string& target_lang = "") const; + std::vector transcribe_16k_nbest( + const std::vector& pcm16k, int beam_size, int nbest, + bool score_norm, const std::string& target_lang) const; + ModelLoader loader_; }; diff --git a/src/parakeet_capi.cpp b/src/parakeet_capi.cpp index 47b0bfe..cfe481b 100644 --- a/src/parakeet_capi.cpp +++ b/src/parakeet_capi.cpp @@ -342,6 +342,87 @@ extern "C" char* parakeet_capi_transcribe_pcm_batch_json(parakeet_ctx* ctx, nullptr); } +extern "C" 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) { + if (!ctx) return nullptr; + if (!ctx->model) { + ctx->last_error = "context has no loaded model"; + return nullptr; + } + if (!wav_path) { + ctx->last_error = "wav_path is NULL"; + return nullptr; + } + try { + const bool normalize = score_norm != 0; + const std::string lang = target_lang ? target_lang : ""; + std::vector hypotheses = + ctx->model->transcribe_path_nbest( + wav_path, beam_size, nbest, normalize, lang); + const pk::ParakeetConfig& cfg = ctx->model->config(); + const float frame_sec = + (float)cfg.hop_length * (float)cfg.subsampling_factor / + (float)cfg.sample_rate; + std::string json = pk::nbest_transcriptions_to_json( + hypotheses, beam_size, normalize, frame_sec); + ctx->last_error.clear(); + char* out = dup_to_c(json); + if (!out) { + ctx->last_error = "out of memory"; + return nullptr; + } + return out; + } catch (const std::exception& e) { + ctx->last_error = e.what(); + return nullptr; + } catch (...) { + ctx->last_error = "unknown error"; + return nullptr; + } +} + +extern "C" 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) { + if (!ctx) return nullptr; + if (!ctx->model) { + ctx->last_error = "context has no loaded model"; + return nullptr; + } + if (!samples || n_samples < 0) { + ctx->last_error = "invalid samples buffer"; + return nullptr; + } + try { + const bool normalize = score_norm != 0; + const std::string lang = target_lang ? target_lang : ""; + const std::vector pcm(samples, samples + n_samples); + std::vector hypotheses = + ctx->model->transcribe_pcm_nbest( + pcm, sample_rate, beam_size, nbest, normalize, lang); + const pk::ParakeetConfig& cfg = ctx->model->config(); + const float frame_sec = + (float)cfg.hop_length * (float)cfg.subsampling_factor / + (float)cfg.sample_rate; + std::string json = pk::nbest_transcriptions_to_json( + hypotheses, beam_size, normalize, frame_sec); + ctx->last_error.clear(); + char* out = dup_to_c(json); + if (!out) { + ctx->last_error = "out of memory"; + return nullptr; + } + return out; + } catch (const std::exception& e) { + ctx->last_error = e.what(); + return nullptr; + } catch (...) { + ctx->last_error = "unknown error"; + return nullptr; + } +} + // --------------------------------------------------------------------------- // Streaming API // --------------------------------------------------------------------------- diff --git a/src/tdt.cpp b/src/tdt.cpp index 72b5f71..31ebc0f 100644 --- a/src/tdt.cpp +++ b/src/tdt.cpp @@ -1,7 +1,10 @@ #include "tdt.hpp" #include "decode_common.hpp" +#include #include #include +#include +#include namespace pk { @@ -117,4 +120,339 @@ std::vector tdt_greedy(const PredictionNet& pred, const Joint& joint, return hyp; } +namespace detail { + +int best_positive_duration_index( + const std::vector& scores, + const std::vector& durations) { + if (scores.size() != durations.size()) + throw std::invalid_argument( + "best_positive_duration_index: size mismatch"); + + int best = -1; + for (size_t i = 0; i < durations.size(); ++i) { + if (durations[i] > 0 && + (best < 0 || scores[i] > scores[(size_t)best])) { + best = (int)i; + } + } + return best; +} + +} // namespace detail + +namespace { + +struct BeamState { + TdtBeamHypothesis hyp; + PredState dec_state; + PredState next_state; + std::vector pred_out; + int last_frame = 0; + int32_t last_token = -1; + bool emitted_any = false; + bool pred_valid = false; +}; + +struct RankedIndex { + float score; + int index; +}; + +std::vector log_softmax(const float* logits, int n) { + float max_logit = logits[0]; + for (int i = 1; i < n; ++i) + max_logit = std::max(max_logit, logits[i]); + + double sum = 0.0; + for (int i = 0; i < n; ++i) + sum += std::exp((double)logits[i] - (double)max_logit); + const double log_denom = std::log(sum); + + std::vector out(n); + for (int i = 0; i < n; ++i) + out[i] = (float)((double)logits[i] - (double)max_logit - log_denom); + return out; +} + +std::vector top_k(const float* scores, int n, int k) { + std::vector ranked; + ranked.reserve(n); + for (int i = 0; i < n; ++i) + ranked.push_back(RankedIndex{scores[i], i}); + + const auto better = [](const RankedIndex& a, const RankedIndex& b) { + if (a.score != b.score) return a.score > b.score; + return a.index < b.index; + }; + const int kept = std::min(n, k); + std::partial_sort(ranked.begin(), ranked.begin() + kept, ranked.end(), better); + ranked.resize(kept); + return ranked; +} + +float log_add_exp(float a, float b) { + const float hi = std::max(a, b); + const float lo = std::min(a, b); + return hi + std::log1p(std::exp(lo - hi)); +} + +bool same_path(const BeamState& a, const BeamState& b) { + if (a.last_frame != b.last_frame || a.hyp.tokens.size() != b.hyp.tokens.size()) + return false; + for (size_t i = 0; i < a.hyp.tokens.size(); ++i) + if (a.hyp.tokens[i].id != b.hyp.tokens[i].id) + return false; + return true; +} + +void merge_duplicate_hypotheses(std::vector& hypotheses) { + std::sort(hypotheses.begin(), hypotheses.end(), + [](const BeamState& a, const BeamState& b) { + return a.hyp.score > b.hyp.score; + }); + + std::vector merged; + merged.reserve(hypotheses.size()); + for (BeamState& candidate : hypotheses) { + auto duplicate = std::find_if( + merged.begin(), merged.end(), + [&](const BeamState& kept) { return same_path(kept, candidate); }); + if (duplicate == merged.end()) { + merged.push_back(std::move(candidate)); + } else { + duplicate->hyp.score = log_add_exp(duplicate->hyp.score, + candidate.hyp.score); + } + } + hypotheses.swap(merged); +} + +float sort_score(const BeamState& hyp, bool score_norm) { + if (!score_norm) return hyp.hyp.score; + return hyp.hyp.score / (float)(hyp.hyp.tokens.size() + 1); +} + +} // namespace + +std::vector tdt_beam_search( + const PredictionNet& pred, const Joint& joint, + const std::vector& enc, int T, int enc_hidden, + const std::vector& durations, int blank_id, + int beam_size, int nbest, bool score_norm) { + if (T < 0 || enc_hidden <= 0 || + enc.size() != (size_t)T * (size_t)enc_hidden) + throw std::invalid_argument("tdt_beam_search: invalid encoder shape"); + if (durations.empty()) + throw std::invalid_argument("tdt_beam_search: durations must not be empty"); + if (beam_size < 1 || nbest < 1 || nbest > beam_size) + throw std::invalid_argument( + "tdt_beam_search: require beam_size >= nbest >= 1"); + + const int num_durations = (int)durations.size(); + const int token_count = joint.V_plus() - num_durations; + if (token_count != joint.vocab_size() + 1 || + num_durations != joint.num_durations() || + blank_id != token_count - 1 || blank_id < 1) { + throw std::invalid_argument("tdt_beam_search: incompatible TDT head"); + } + + int zero_duration_idx = -1; + int zero_duration_count = 0; + bool has_positive_duration = false; + for (int i = 0; i < num_durations; ++i) { + if (durations[i] < 0) + throw std::invalid_argument( + "tdt_beam_search: durations must be non-negative"); + if (durations[i] == 0) { + zero_duration_idx = i; + ++zero_duration_count; + } + if (durations[i] > 0) + has_positive_duration = true; + } + if (!has_positive_duration) + throw std::invalid_argument( + "tdt_beam_search: at least one positive duration is required"); + if (zero_duration_count > 1) + throw std::invalid_argument( + "tdt_beam_search: at most one zero duration is allowed"); + if (T == 0) + return {TdtBeamHypothesis{}}; + + const int beam = std::min(beam_size, blank_id); + const int token_beam = std::min(beam, blank_id - 1); + const int duration_beam = std::min(beam, num_durations); + + std::vector enc_proj; + joint.precompute_enc_proj(enc, T, enc_hidden, enc_proj); + const int joint_hidden = joint.joint_hidden(); + + BeamState start; + start.dec_state = pred.zero_state(); + std::vector kept_hyps; + kept_hyps.push_back(std::move(start)); + + std::vector logits; + for (int time_idx = 0; time_idx < T; ++time_idx) { + std::vector current_hyps; + std::vector future_hyps; + current_hyps.reserve(kept_hyps.size()); + future_hyps.reserve(kept_hyps.size()); + for (BeamState& hyp : kept_hyps) { + if (hyp.last_frame == time_idx) + current_hyps.push_back(std::move(hyp)); + else if (hyp.last_frame > time_idx) + future_hyps.push_back(std::move(hyp)); + } + kept_hyps.swap(future_hyps); + + while (!current_hyps.empty()) { + auto best_it = std::max_element( + current_hyps.begin(), current_hyps.end(), + [](const BeamState& a, const BeamState& b) { + return a.hyp.score < b.hyp.score; + }); + BeamState best = std::move(*best_it); + current_hyps.erase(best_it); + + if (!best.pred_valid) { + const bool is_sos = !best.emitted_any; + const int32_t label = best.emitted_any ? best.last_token : blank_id; + pred.step(label, is_sos, best.dec_state, + best.pred_out, best.next_state); + best.pred_valid = true; + } + + joint.step_logits( + enc_proj.data() + (size_t)time_idx * joint_hidden, + best.pred_out.data(), (int)best.pred_out.size(), logits); + std::vector token_logp = + log_softmax(logits.data(), token_count); + std::vector duration_logp = + log_softmax(logits.data() + token_count, num_durations); + if (!std::all_of(token_logp.begin(), token_logp.end(), + [](float value) { return std::isfinite(value); }) || + !std::all_of(duration_logp.begin(), duration_logp.end(), + [](float value) { return std::isfinite(value); })) { + throw std::runtime_error( + "tdt_beam_search: non-finite log probability"); + } + + const std::vector best_durations = + top_k(duration_logp.data(), num_durations, duration_beam); + + struct Pair { + float score; + int token; + int duration_idx; + }; + std::vector pairs; + const std::vector best_tokens = + top_k(token_logp.data(), blank_id, token_beam); + pairs.reserve(best_tokens.size() * best_durations.size()); + for (const RankedIndex& duration : best_durations) + for (const RankedIndex& token : best_tokens) + pairs.push_back(Pair{ + duration.score + token.score, + token.index, + duration.index}); + std::partial_sort( + pairs.begin(), + pairs.begin() + std::min(token_beam, (int)pairs.size()), + pairs.end(), + [](const Pair& a, const Pair& b) { + if (a.score != b.score) return a.score > b.score; + if (a.duration_idx != b.duration_idx) + return a.duration_idx < b.duration_idx; + return a.token < b.token; + }); + pairs.resize(std::min(token_beam, (int)pairs.size())); + + for (const Pair& pair : pairs) { + const int duration = durations[pair.duration_idx]; + BeamState child = best; + child.hyp.score += pair.score; + if (duration == 0 && + !(child.hyp.score < best.hyp.score)) { + throw std::runtime_error( + "tdt_beam_search: zero-duration expansion " + "did not reduce score"); + } + child.hyp.tokens.push_back(TdtBeamToken{ + (int32_t)pair.token, (int32_t)time_idx, (int32_t)duration}); + child.dec_state = best.next_state; + child.last_token = (int32_t)pair.token; + child.emitted_any = true; + child.pred_valid = false; + child.last_frame += duration; + if (duration == 0) + current_hyps.push_back(std::move(child)); + else + kept_hyps.push_back(std::move(child)); + } + + for (const RankedIndex& ranked_duration : best_durations) { + int duration_idx = ranked_duration.index; + if (duration_idx == zero_duration_idx) { + if (best_durations.size() != 1) + continue; + duration_idx = detail::best_positive_duration_index( + duration_logp, durations); + } + BeamState child = best; + child.hyp.score += token_logp[blank_id] + + duration_logp[duration_idx]; + child.last_frame += durations[duration_idx]; + kept_hyps.push_back(std::move(child)); + } + + merge_duplicate_hypotheses(kept_hyps); + if (!current_hyps.empty()) { + const float next_score = std::max_element( + current_hyps.begin(), current_hyps.end(), + [](const BeamState& a, const BeamState& b) { + return a.hyp.score < b.hyp.score; + })->hyp.score; + const int more_probable_count = (int)std::count_if( + kept_hyps.begin(), kept_hyps.end(), + [next_score](const BeamState& hyp) { + return hyp.hyp.score > next_score; + }); + if (more_probable_count >= beam) { + std::vector more_probable; + more_probable.reserve(more_probable_count); + for (BeamState& hyp : kept_hyps) + if (hyp.hyp.score > next_score) + more_probable.push_back(std::move(hyp)); + kept_hyps.swap(more_probable); + break; + } + } else { + std::sort(kept_hyps.begin(), kept_hyps.end(), + [](const BeamState& a, const BeamState& b) { + return a.hyp.score > b.hyp.score; + }); + if ((int)kept_hyps.size() > beam) + kept_hyps.resize(beam); + } + } + } + + std::sort(kept_hyps.begin(), kept_hyps.end(), + [score_norm](const BeamState& a, const BeamState& b) { + return sort_score(a, score_norm) > sort_score(b, score_norm); + }); + const int count = std::min(nbest, (int)kept_hyps.size()); + std::vector result; + result.reserve(count); + for (int i = 0; i < count; ++i) { + kept_hyps[i].hyp.normalized_score = + kept_hyps[i].hyp.score / + (float)(kept_hyps[i].hyp.tokens.size() + 1); + result.push_back(std::move(kept_hyps[i].hyp)); + } + return result; +} + } // namespace pk diff --git a/src/tdt.hpp b/src/tdt.hpp index 73eb738..66acb5b 100644 --- a/src/tdt.hpp +++ b/src/tdt.hpp @@ -7,6 +7,27 @@ namespace pk { +// One non-blank token emitted by TDT beam search. +// +// `frame` is the encoder frame at which the token was emitted. `duration` is +// the selected TDT duration, so the token's end frame is frame + duration. +struct TdtBeamToken { + int32_t id; + int32_t frame; + int32_t duration; +}; + +// One completed TDT beam-search hypothesis. +// +// `score` is the accumulated token+duration log probability. NeMo's default +// score normalization divides by the sequence length INCLUDING its leading +// blank/SOS sentinel, hence tokens.size() + 1 here. +struct TdtBeamHypothesis { + std::vector tokens; + float score = 0.0f; + float normalized_score = 0.0f; +}; + // TDT (Token-and-Duration Transducer) duration-aware greedy decoding. // // Ports NeMo GreedyTDTInfer._greedy_decode (rnnt_greedy_decoding.py). Drives the @@ -56,4 +77,37 @@ std::vector tdt_greedy(const PredictionNet& pred, const Joint& joint, int blank_id, int max_symbols, std::vector* tokens = nullptr); +// Sequence-level beam search for TDT models, matching NeMo BeamTDTInfer's +// `default_beam_search`: +// +// * token and duration slices are log-softmaxed independently; +// * the best non-blank token-duration pairs are expanded jointly; +// * blank is expanded only with non-zero durations; +// * duplicate (token sequence, last frame) paths are merged with logaddexp; +// * final hypotheses are sorted by normalized score by default. +// +// A zero-duration expansion must strictly reduce the accumulated log score. +// Rejecting non-finite or non-decreasing scores prevents a malformed model from +// creating a non-progress loop without truncating valid hypotheses. +// +// This is an opt-in offline decoder. The existing tdt_greedy path is not used +// or modified. `beam_size` and `nbest` must be positive, with nbest <= +// beam_size. At most `nbest` hypotheses are returned. +std::vector tdt_beam_search( + const PredictionNet& pred, const Joint& joint, + const std::vector& enc, int T, int enc_hidden, + const std::vector& durations, int blank_id, + int beam_size, int nbest, bool score_norm = true); + +namespace detail { + +// Returns the index of the highest-scoring positive duration, or -1 when no +// positive duration exists. Kept in the internal TDT header so the beam-size-1 +// blank fallback can be tested without loading a model. +int best_positive_duration_index( + const std::vector& scores, + const std::vector& durations); + +} // namespace detail + } // namespace pk diff --git a/src/transcription_json.cpp b/src/transcription_json.cpp index 130b0ae..7d81de5 100644 --- a/src/transcription_json.cpp +++ b/src/transcription_json.cpp @@ -1,4 +1,5 @@ #include "transcription_json.hpp" +#include "model.hpp" #include @@ -91,4 +92,47 @@ std::string transcription_to_json(const Transcription& tr, float frame_sec) { return out; } +std::string nbest_transcriptions_to_json( + const std::vector& hypotheses, + int beam_size, bool score_norm, float frame_sec) { + std::string out; + out.reserve(96 + hypotheses.size() * 160); + out += "{\"beam_size\":"; + append_json_int(out, beam_size); + out += ",\"score_norm\":"; + out += score_norm ? "true" : "false"; + out += ",\"frame_sec\":"; + append_json_float(out, "%.6f", frame_sec); + out += ",\"hypotheses\":["; + for (size_t i = 0; i < hypotheses.size(); ++i) { + if (i) out += ','; + const NBestTranscription& hyp = hypotheses[i]; + out += "{\"text\":"; + append_json_string(out, hyp.text); + out += ",\"score\":"; + append_json_float(out, "%.6f", hyp.score); + out += ",\"normalized_score\":"; + append_json_float(out, "%.6f", hyp.normalized_score); + out += ",\"tokens\":["; + for (size_t j = 0; j < hyp.tokens.size(); ++j) { + if (j) out += ','; + const TdtBeamToken& token = hyp.tokens[j]; + out += "{\"id\":"; + append_json_int(out, token.id); + out += ",\"frame\":"; + append_json_int(out, token.frame); + out += ",\"t\":"; + append_json_float(out, "%.3f", (float)token.frame * frame_sec); + out += ",\"duration_frames\":"; + append_json_int(out, token.duration); + out += ",\"duration\":"; + append_json_float(out, "%.3f", (float)token.duration * frame_sec); + out += '}'; + } + out += "]}"; + } + out += "]}"; + return out; +} + } // namespace pk diff --git a/src/transcription_json.hpp b/src/transcription_json.hpp index bf4d6c2..1bda990 100644 --- a/src/transcription_json.hpp +++ b/src/transcription_json.hpp @@ -3,9 +3,12 @@ #include "transcription.hpp" #include +#include namespace pk { +struct NBestTranscription; + void append_json_string(std::string& out, const std::string& s); void append_json_int(std::string& out, int v); void append_json_float(std::string& out, const char* fmt, float v); @@ -14,4 +17,10 @@ void append_json_float(std::string& out, const char* fmt, float v); // {"text", "frame_sec", "words", "tokens"}. std::string transcription_to_json(const Transcription& tr, float frame_sec); +// Serialize the offline TDT beam result: +// {"beam_size","score_norm","frame_sec","hypotheses":[...]}. +std::string nbest_transcriptions_to_json( + const std::vector& hypotheses, + int beam_size, bool score_norm, float frame_sec); + } // namespace pk diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8c9c175..126b48e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -7,6 +7,8 @@ function(pk_add_test name) endfunction() pk_add_test(test_smoke) +pk_add_test(test_nbest_json) +pk_add_test(test_tdt_beam_core) pk_add_test(test_backend_device) pk_add_test(test_audio_io) pk_add_test(test_model_loader) @@ -40,6 +42,7 @@ pk_add_test(test_joint_step_batch) pk_add_test(test_prompt_kernel) pk_add_test(test_transducer_core) pk_add_test(test_tdt_greedy) +pk_add_test(test_tdt_beam) pk_add_test(test_transducer_greedy_batch) pk_add_test(test_transducer_greedy_batch_rnnt) pk_add_test(test_timestamps_tokens) @@ -116,6 +119,7 @@ set_tests_properties(test_model_loader test_mel test_mel_gpu test_subsampling te test_streaming_encoder test_ctc test_prediction test_prediction_step test_prediction_step_batch test_joint test_joint_step_batch test_prompt_kernel test_transducer_core test_tdt_greedy + test_tdt_beam test_transducer_greedy_batch test_transducer_greedy_batch_rnnt test_timestamps_tokens test_timestamps test_transcribe_batch_ts test_tokenizer test_transcribe test_transcribe_speech test_transcribe_tiled test_transcribe_tdt test_transcribe_0_6b @@ -129,6 +133,7 @@ set_tests_properties(test_mel test_mel_gpu test_subsampling test_subsampling_bat test_ctc test_prediction test_prediction_step test_prediction_step_batch test_joint test_joint_step_batch test_prompt_kernel test_transducer_core test_tdt_greedy + test_tdt_beam test_transducer_greedy_batch test_transducer_greedy_batch_rnnt test_timestamps_tokens test_timestamps test_transcribe_batch_ts diff --git a/tests/test_nbest_json.cpp b/tests/test_nbest_json.cpp new file mode 100644 index 0000000..53c34b7 --- /dev/null +++ b/tests/test_nbest_json.cpp @@ -0,0 +1,28 @@ +#include "model.hpp" +#include "transcription_json.hpp" + +#include +#include +#include + +int main() { + pk::NBestTranscription first; + first.text = "hello"; + first.score = -2.0f; + first.normalized_score = -1.0f; + first.tokens.push_back(pk::TdtBeamToken{3, 2, 4}); + + const std::string got = pk::nbest_transcriptions_to_json( + std::vector{first}, 4, true, 0.08f); + const std::string expected = + "{\"beam_size\":4,\"score_norm\":true,\"frame_sec\":0.080000," + "\"hypotheses\":[{\"text\":\"hello\",\"score\":-2.000000," + "\"normalized_score\":-1.000000,\"tokens\":[{\"id\":3,\"frame\":2," + "\"t\":0.160,\"duration_frames\":4,\"duration\":0.320}]}]}"; + if (got != expected) { + std::fprintf(stderr, "test_nbest_json: mismatch\n got: %s\n exp: %s\n", + got.c_str(), expected.c_str()); + return 1; + } + return 0; +} diff --git a/tests/test_tdt_beam.cpp b/tests/test_tdt_beam.cpp new file mode 100644 index 0000000..8a17ef1 --- /dev/null +++ b/tests/test_tdt_beam.cpp @@ -0,0 +1,310 @@ +#include "audio_io.hpp" +#include "ggml_graph.hpp" +#include "model.hpp" +#include "parakeet_capi.h" +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct BackendShutdown { + ~BackendShutdown() { pk::shutdown_backend(); } +}; + +} // namespace + +static bool read_i32_tensor(ggml_context* ctx, const char* name, + std::vector& out) { + ggml_tensor* tensor = ggml_get_tensor(ctx, name); + if (!tensor || tensor->type != GGML_TYPE_I32) + return false; + const int n = (int)ggml_nelements(tensor); + const int32_t* data = static_cast(tensor->data); + out.assign(data, data + n); + return true; +} + +static bool read_f32_tensor(ggml_context* ctx, const char* name, + std::vector& out) { + ggml_tensor* tensor = ggml_get_tensor(ctx, name); + if (!tensor || tensor->type != GGML_TYPE_F32) + return false; + const int n = (int)ggml_nelements(tensor); + const float* data = static_cast(tensor->data); + out.assign(data, data + n); + return true; +} + +static bool compare_nemo_baseline( + const char* path, int beam_size, + const std::vector& got) { + ggml_context* ctx = nullptr; + gguf_init_params params{/*no_alloc=*/false, &ctx}; + gguf_context* gguf = gguf_init_from_file(path, params); + if (!gguf) { + std::fprintf(stderr, + "test_tdt_beam: failed to open baseline %s\n", path); + return false; + } + + const int64_t beam_key = + gguf_find_key(gguf, "baseline.tdt_nbest_beam_size"); + const int64_t count_key = + gguf_find_key(gguf, "baseline.tdt_nbest_count"); + if (beam_key < 0 || count_key < 0) { + std::fprintf(stderr, + "test_tdt_beam: baseline has no TDT N-best data\n"); + gguf_free(gguf); + ggml_free(ctx); + return false; + } + const int reference_beam = (int)gguf_get_val_u32(gguf, beam_key); + const int reference_count = (int)gguf_get_val_u32(gguf, count_key); + + std::vector offsets; + std::vector ids; + std::vector end_frames; + std::vector scores; + const bool tensors_ok = + read_i32_tensor(ctx, "tdt_nbest_offsets", offsets) && + read_i32_tensor(ctx, "tdt_nbest_token_ids", ids) && + read_i32_tensor(ctx, "tdt_nbest_token_end_frames", end_frames) && + read_f32_tensor(ctx, "tdt_nbest_scores", scores); + gguf_free(gguf); + ggml_free(ctx); + if (!tensors_ok || reference_beam != beam_size || + reference_count != (int)got.size() || + offsets.size() != got.size() + 1 || + scores.size() != got.size() || + ids.size() != end_frames.size()) { + std::fprintf(stderr, + "test_tdt_beam: invalid N-best baseline shape\n"); + return false; + } + + for (size_t rank = 0; rank < got.size(); ++rank) { + const int begin = offsets[rank]; + const int end = offsets[rank + 1]; + if (begin < 0 || end < begin || end > (int)ids.size() || + end - begin != (int)got[rank].tokens.size()) { + std::fprintf(stderr, + "test_tdt_beam: rank %zu token count mismatch\n", + rank); + return false; + } + if (std::fabs(scores[rank] - got[rank].score) > 5e-3f) { + std::fprintf(stderr, + "test_tdt_beam: rank %zu score mismatch " + "got=%.6f ref=%.6f\n", + rank, got[rank].score, scores[rank]); + return false; + } + for (int i = begin; i < end; ++i) { + const pk::TdtBeamToken& token = + got[rank].tokens[(size_t)(i - begin)]; + if (token.id != ids[i] || + token.frame + token.duration != end_frames[i]) { + std::fprintf(stderr, + "test_tdt_beam: rank %zu token %d mismatch " + "got=(%d,%d) ref=(%d,%d)\n", + rank, i - begin, token.id, + token.frame + token.duration, + ids[i], end_frames[i]); + return false; + } + } + } + return true; +} + +static bool same_hypothesis(const pk::NBestTranscription& a, + const pk::NBestTranscription& b) { + if (a.text != b.text || a.tokens.size() != b.tokens.size()) + return false; + if (std::fabs(a.score - b.score) > 1e-6f || + std::fabs(a.normalized_score - b.normalized_score) > 1e-6f) + return false; + for (size_t i = 0; i < a.tokens.size(); ++i) { + if (a.tokens[i].id != b.tokens[i].id || + a.tokens[i].frame != b.tokens[i].frame || + a.tokens[i].duration != b.tokens[i].duration) + return false; + } + return true; +} + +int main() { + BackendShutdown backend_shutdown; + const char* model_path = std::getenv("PARAKEET_TEST_GGUF"); + if (!model_path) { + std::fprintf(stderr, + "test_tdt_beam: PARAKEET_TEST_GGUF not set; skip\n"); + return 77; + } + + std::unique_ptr model = pk::Model::load(model_path); + if (!model) { + std::fprintf(stderr, "test_tdt_beam: failed to load model\n"); + return 1; + } + const pk::ParakeetConfig& cfg = model->config(); + if (cfg.tdt_durations.empty()) { + std::fprintf(stderr, "test_tdt_beam: model is not TDT; skip\n"); + return 77; + } + + const int beam_size = 4; + std::vector first = + model->transcribe_path_nbest( + "tests/fixtures/speech.wav", beam_size, beam_size, true); + if (first.empty() || first.size() > (size_t)beam_size) { + std::fprintf(stderr, + "test_tdt_beam: invalid hypothesis count %zu\n", + first.size()); + return 1; + } + + for (size_t i = 0; i < first.size(); ++i) { + const pk::NBestTranscription& hyp = first[i]; + const float expected_norm = + hyp.score / (float)(hyp.tokens.size() + 1); + if (!std::isfinite(hyp.score) || + std::fabs(hyp.normalized_score - expected_norm) > 1e-6f) { + std::fprintf(stderr, + "test_tdt_beam: invalid score at rank %zu\n", i); + return 1; + } + if (i > 0 && + first[i - 1].normalized_score < hyp.normalized_score) { + std::fprintf(stderr, + "test_tdt_beam: hypotheses are not score-sorted\n"); + return 1; + } + int previous_frame = -1; + for (const pk::TdtBeamToken& token : hyp.tokens) { + if (token.id < 0 || token.id >= (int)cfg.blank_id || + token.frame < previous_frame || + std::find(cfg.tdt_durations.begin(), cfg.tdt_durations.end(), + token.duration) == cfg.tdt_durations.end()) { + std::fprintf(stderr, + "test_tdt_beam: invalid token metadata\n"); + return 1; + } + previous_frame = token.frame; + } + } + + std::vector second = + model->transcribe_path_nbest( + "tests/fixtures/speech.wav", beam_size, beam_size, true); + if (first.size() != second.size()) { + std::fprintf(stderr, + "test_tdt_beam: nondeterministic hypothesis count\n"); + return 1; + } + for (size_t i = 0; i < first.size(); ++i) { + if (!same_hypothesis(first[i], second[i])) { + std::fprintf(stderr, + "test_tdt_beam: nondeterministic rank %zu\n", i); + return 1; + } + } + + std::vector raw_ranked = + model->transcribe_path_nbest( + "tests/fixtures/speech.wav", beam_size, 2, false); + if (raw_ranked.empty() || raw_ranked.size() > 2) { + std::fprintf(stderr, + "test_tdt_beam: invalid raw-ranked hypothesis count\n"); + return 1; + } + for (size_t i = 1; i < raw_ranked.size(); ++i) { + if (raw_ranked[i - 1].score < raw_ranked[i].score) { + std::fprintf(stderr, + "test_tdt_beam: hypotheses are not raw-score sorted\n"); + return 1; + } + } + + try { + (void)model->transcribe_path_nbest( + "tests/fixtures/does-not-exist.wav", 0, 1, true); + std::fprintf(stderr, + "test_tdt_beam: invalid beam request unexpectedly succeeded\n"); + return 1; + } catch (const std::invalid_argument& e) { + if (!std::strstr(e.what(), "beam_size")) { + std::fprintf(stderr, + "test_tdt_beam: unexpected validation error: %s\n", + e.what()); + return 1; + } + } + + const char* baseline_path = + std::getenv("PARAKEET_TEST_TDT_NBEST_BASELINE"); + if (baseline_path && + !compare_nemo_baseline(baseline_path, beam_size, first)) { + return 1; + } + + model.reset(); + parakeet_ctx* ctx = parakeet_capi_load(model_path); + if (!ctx) { + std::fprintf(stderr, "test_tdt_beam: C API load failed\n"); + return 1; + } + char* json = parakeet_capi_transcribe_path_nbest_json( + ctx, "tests/fixtures/speech.wav", + beam_size, beam_size, 1, nullptr); + if (!json || !std::strstr(json, "\"hypotheses\":[") || + !std::strstr(json, "\"normalized_score\":")) { + std::fprintf(stderr, "test_tdt_beam: invalid C API JSON: %s\n", + json ? json : parakeet_capi_last_error(ctx)); + parakeet_capi_free_string(json); + parakeet_capi_free(ctx); + return 1; + } + parakeet_capi_free_string(json); + + pk::Audio audio; + if (!pk::load_audio_16k_mono("tests/fixtures/speech.wav", audio)) { + std::fprintf(stderr, "test_tdt_beam: failed to load PCM fixture\n"); + parakeet_capi_free(ctx); + return 1; + } + char* pcm_json = parakeet_capi_transcribe_pcm_nbest_json( + ctx, audio.samples.data(), (int)audio.samples.size(), audio.sample_rate, + beam_size, 2, 0, nullptr); + if (!pcm_json || + !std::strstr(pcm_json, "\"score_norm\":false") || + !std::strstr(pcm_json, "\"hypotheses\":[")) { + std::fprintf(stderr, "test_tdt_beam: invalid PCM C API JSON: %s\n", + pcm_json ? pcm_json : parakeet_capi_last_error(ctx)); + parakeet_capi_free_string(pcm_json); + parakeet_capi_free(ctx); + return 1; + } + parakeet_capi_free_string(pcm_json); + + char* invalid = parakeet_capi_transcribe_path_nbest_json( + ctx, "tests/fixtures/does-not-exist.wav", 0, 1, 1, nullptr); + if (invalid || !std::strstr(parakeet_capi_last_error(ctx), "beam_size")) { + std::fprintf(stderr, + "test_tdt_beam: C API validation did not fail early\n"); + parakeet_capi_free_string(invalid); + parakeet_capi_free(ctx); + return 1; + } + parakeet_capi_free(ctx); + return 0; +} diff --git a/tests/test_tdt_beam_core.cpp b/tests/test_tdt_beam_core.cpp new file mode 100644 index 0000000..5d46026 --- /dev/null +++ b/tests/test_tdt_beam_core.cpp @@ -0,0 +1,20 @@ +#include "tdt.hpp" + +#include +#include + +int main() { + const std::vector durations{0, 1, 2, 4}; + const std::vector scores{0.0f, -3.0f, -0.1f, -2.0f}; + + const int best = + pk::detail::best_positive_duration_index(scores, durations); + if (best != 2) { + std::fprintf( + stderr, + "test_tdt_beam_core: expected duration index 2, got %d\n", + best); + return 1; + } + return 0; +}