diff --git a/AGENTS.md b/AGENTS.md index c650ce3..830bf0b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -52,6 +52,9 @@ before, so do not change them without an A/B benchmark that proves parity. (so the unsupported op can run on CPU); when every op is supported, the fast gallocr path runs. If you think gallocr can go, you are about to reintroduce that regression. +- **Ternary weights stay packed in the GGUF and are repacked once per loader; do not dequantize per call.** + The packed Redux form is what makes it 6.8x smaller than F16 and what the + CPU kernels in `src/ternary*.cpp` read. - **Zero-copy weights.** `clone_weight` returns loader tensors directly so the same device buffer is reused every utterance; do not copy weights per call. @@ -84,6 +87,10 @@ src/ libparakeet implementation ced_tagger.hpp/cpp , pk::CedTagger: loads a CED GGUF (ced.cpp) into a tagger context, pk::SoundScorer sound_stream.hpp/cpp, pk::SoundStream: sliding-window sound-event detection over live PCM scene_stream.hpp/cpp, pk::SceneStream: combined ASR + diarization + sound-event stream + ternary.hpp/cpp , packed ternary (moondream/parakeet-redux) linears: repack once, int8 activations, scalar ref + two-op ggml custom op + ternary_kernels.hpp, ternary_kernels_x86.cpp (AVX-512 VNNI / AVX2), ternary_kernels_neon.cpp + vad_head.hpp/cpp , voice-activity head of Ultra/Redux (plain C++ loops), Model::vad_probabilities + vad_segmenter.hpp/cpp, pk::segment_by_vad: cut long audio at VAD pauses (SegmenterOpts), used by `transcribe --vad` scene_render.hpp/cpp, pk::SceneRenderer + format_span/is_speech_label: `parakeet-cli scene` text rendering examples/cli/ parakeet-cli binary subcommands: info, transcribe (+ --stream), quantize, scene (ASR + diar + sound, one time-ordered feed) @@ -91,6 +98,7 @@ examples/cli/ parakeet-cli binary diarize binary: diarize [--stream] scripts/ Python tooling convert_parakeet_to_gguf.py, .nemo/.hf -> GGUF (--dtype f32|f16|q8_0) + convert_hf_parakeet_to_gguf.py, HF safetensors (moondream/parakeet-ultra, -redux) to GGUF (--template, --ternary keep|dequant, --vad keep|drop) gen_nemo_baseline.py , NeMo intermediates -> baseline.gguf gen_stream_baseline.py , NeMo cache-aware streaming encode+decode -> stream baseline.gguf gen_diar_baseline.py , NeMo offline + streaming diarization -> diar baseline.gguf @@ -117,6 +125,13 @@ tests/ ctest targets test_streaming_diarization.cpp, streaming diarization == NeMo streaming, every latency mode (same baseline) test_combined_offline.cpp, SAS + streaming diarization/SAS through the C-API test_sas_merge.cpp , SAS merge/grouping (model-independent) + test_ternary.cpp , ternary repack, int8 quant, every kernel == scalar (model-independent) + test_ternary_model.cpp , packed Redux == dequantized Redux transcript (PARAKEET_TEST_GGUF_REDUX_KEEP + _DEQ) + test_model_loader_ternary.cpp, ternary + VAD flags from GGUF KVs + bench_ternary.cpp , single-thread throughput of each ternary kernel (not a ctest) + test_vad_head.cpp , VAD head probabilities (PARAKEET_TEST_GGUF_ULTRA) + test_vad_segmenter.cpp , segmenter cut rules (model-independent) + test_transcribe_vad.cpp , --vad path vs plain pass on long audio (PARAKEET_TEST_GGUF_ULTRA, PARAKEET_TEST_GGUF) test_asr_committer.cpp , shared word/utterance finalize logic (model-independent) test_ced_parity.cpp , CedTagger scores == ced.cpp PyTorch baseline (PARAKEET_TEST_CED_GGUF f32 + PARAKEET_TEST_CED_BASELINE) test_sound_stream.cpp , pk::SoundStream windowing/on-off-min_duration logic (model-independent) @@ -141,6 +156,7 @@ docs/ conversion.md , GGUF schema reference quantization.md , quantization allowlist, policy, measured size + WER per type parity.md , full model coverage matrix + per-stage tensor parity + ternary.md , packed ternary Redux: GGUF form, kernels, limits, measured speed diarization.md , speaker diarization + speaker-attributed ASR: parity, C-API, speed .github/workflows/ ci.yml , build job (per-push) + closed-loop job (pull_request + dispatch) @@ -189,6 +205,10 @@ ctest --test-dir build --output-on-failure Tests return exit code 77 (ctest SKIP) when the venv or checkpoint is absent, so they never break a CI environment that lacks them. +Ultra/Redux tests read `PARAKEET_TEST_GGUF_ULTRA` (F16 Ultra), +`PARAKEET_TEST_GGUF_REDUX_KEEP` (packed ternary) and `PARAKEET_TEST_GGUF_REDUX_DEQ` +(dequantized Redux); they skip (77) when unset. + ### Test labels | Label | Tests | Needs | @@ -221,6 +241,17 @@ Convert (HuggingFace id or local `.nemo`): Featurizer window and filterbank are lifted from the checkpoint at runtime; mel/fft parameters do not need to be specified manually. +## Ternary GGUF flags + +Two optional GGUF flags, both read into `ParakeetConfig`: + +- `parakeet.ternary.present` (with `parakeet.ternary.group_size` = 128): the + encoder linears are stored as `.qweight` (I8) + `.scales` (F16). + CPU only, offline only (no streaming). `PARAKEET_TERNARY_KERNEL=scalar|avx2|vnni|neon` + forces a kernel. See `docs/ternary.md`. +- `parakeet.vad.present` (with `parakeet.vad.d_in/hidden/kernel/frame_sec`): the + file carries `vad_head.*` tensors. + ## Quantization policy See `docs/quantization.md` for the full policy. Summary: @@ -255,6 +286,8 @@ The binary is at `build/examples/cli/parakeet-cli`. parakeet-cli info parakeet-cli transcribe --model --input [--decoder ctc|tdt] [--stream] [--timestamps] [--json] parakeet-cli quantize +parakeet-cli transcribe --model --input --vad [--vad-threshold F] [--vad-min-pause SEC] [--vad-max-seg SEC] +parakeet-cli vad-probe --model --input [--variant N] # dump VAD head probabilities as t_sec,p parakeet-cli scene [--model ] [--diar ] [--sound ] --input [--latency model|low|very_low|ultra_low] [--chunk-ms N] [--show-speech] [--json] ``` @@ -288,6 +321,13 @@ parakeet_capi_stream_finalize # flush the end-of-stream tail parakeet_capi_stream_free ``` +VAD segmentation (additive, ABI unchanged; not used by LocalAI yet). Needs a GGUF +with a VAD head (Ultra/Redux); see `docs/ternary.md`: + +``` +parakeet_capi_transcribe_path_json_vad # same JSON as _json, long audio cut at VAD pauses +``` + Speaker diarization (ABI v7, additive; not used by LocalAI yet). A diarization GGUF loads into its own `parakeet_ctx`; see `docs/diarization.md`: diff --git a/CMakeLists.txt b/CMakeLists.txt index cbaa2db..d375d4b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -94,6 +94,8 @@ endif() set(PARAKEET_SRC src/parakeet.cpp src/model.cpp + src/vad_head.cpp + src/vad_segmenter.cpp src/parakeet_capi.cpp src/common.cpp src/audio_io.cpp @@ -115,6 +117,9 @@ set(PARAKEET_SRC src/prediction.cpp src/joint.cpp src/prompt_kernel.cpp + src/ternary.cpp + src/ternary_kernels_x86.cpp + src/ternary_kernels_neon.cpp src/tdt.cpp src/rnnt.cpp src/transducer_batch.cpp @@ -143,6 +148,13 @@ if(PARAKEET_SHARED) else() add_library(parakeet STATIC ${PARAKEET_SRC}) endif() +if(NOT MSVC) + set_source_files_properties(src/ternary.cpp src/ternary_kernels_x86.cpp src/ternary_kernels_neon.cpp + PROPERTIES COMPILE_OPTIONS "-ffp-contract=off") +endif() +if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64" AND NOT APPLE AND NOT MSVC) + set_property(SOURCE src/ternary_kernels_neon.cpp APPEND PROPERTY COMPILE_OPTIONS "-march=armv8.2-a+dotprod") +endif() target_include_directories(parakeet PUBLIC include PRIVATE src ${CMAKE_SOURCE_DIR}/third_party) target_compile_definitions(parakeet PUBLIC $<$:_USE_MATH_DEFINES>) target_compile_definitions(parakeet PRIVATE PARAKEET_VERSION="${PARAKEET_VERSION}") diff --git a/README.md b/README.md index 0a3503e..97f5d66 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,26 @@ Every model below is validated at WER 0 against NeMo and published as GGUF (f16, | [parakeet_realtime_eou_120m-v1](https://huggingface.co/nvidia/parakeet_realtime_eou_120m-v1) | RNNT, streaming | 120M | cache-aware streaming with end-of-utterance detection (`--stream`) | NVIDIA | | [nemotron-3.5-asr-streaming-0.6b](https://huggingface.co/nvidia/nemotron-3.5-asr-streaming-0.6b) | RNNT, streaming | 0.6B | multilingual (40+ locales), prompt-conditioned, offline and cache-aware streaming, pick a language with `--lang` (default `auto`). OpenMDW-1.1 | NVIDIA | + +### Moondream Ultra and Redux (not yet published) + +[moondream/parakeet-ultra](https://huggingface.co/moondream/parakeet-ultra) and +[moondream/parakeet-redux](https://huggingface.co/moondream/parakeet-redux) are Moondream's +post-trained (Ultra, F16) and ternary-encoder (Redux) derivatives of parakeet-tdt-0.6b-v3. They are +HF safetensors, converted with `scripts/convert_hf_parakeet_to_gguf.py`. They are not part of the +NeMo-validated set above: there is no NeMo baseline for them, so parity is transcript-level against +our own v3 path (see [`docs/parity.md`](docs/parity.md)), and no GGUFs are published yet. + +- Redux packs the encoder as ternary weights: a 213 MB GGUF, 6.8x smaller than F16. It runs on CPU + only and offline only. On x86 with AVX-512 VNNI it reaches median RTF 75.6 per utterance on + LibriSpeech-100 (8 threads) against 46.1 for the same model in F16; on a single 180 s clip the + gain is about 10 percent; WER on the 100 LibriSpeech utterances is 1.96 percent. See + [`docs/ternary.md`](docs/ternary.md). SIMD kernels exist for x86-64 with AVX2 or AVX-512 VNNI and + aarch64 with dotprod; MSVC builds, Windows on ARM and aarch64 without dotprod use a slow scalar + kernel (about 1 GMAC/s), and the load logs a warning. The packed file also stays resident next to + the repacked planes, so memory use is more than the file size. +- Both carry a voice-activity head, used by `transcribe --vad` to cut long audio at pauses. On + synthetic long-form clips it does not change WER meaningfully. --- ## Performance @@ -242,6 +262,10 @@ parakeet-cli transcribe --model m.gguf --input audio.wav --decoder tdt \ # Read WAV bytes from stdin (useful with ffmpeg/curl pipelines) ffmpeg -i input.mp3 -f wav - | parakeet-cli transcribe --model m.gguf --input - +# Long audio on Ultra/Redux: cut at VAD pauses, transcribe each piece (offline only). +# Tune with --vad-threshold F, --vad-min-pause SEC, --vad-max-seg SEC +parakeet-cli transcribe --model ultra.gguf --input long.wav --vad + # Print model metadata (arch, dims, mel params, vocab size, TDT durations) parakeet-cli info m.gguf @@ -299,7 +323,7 @@ OpenAI API surface, auth, and metrics. Single-clip transcription is the default and needs no flags: every `transcribe` call runs one clip at a time, byte-for-byte identical to before. Batching is an opt-in path for decoding several clips together, which matters when you serve many concurrent requests on a GPU. -The win is on the **decode** side. A transducer (TDT/RNN-T) decodes autoregressively with tiny per-step prediction-LSTM and joint GEMMs; one clip launches hundreds of these matvec-sized kernels and leaves the GPU mostly idle between launches. Decoding N clips together coalesces each step into one batched GEMM, so the device stays busy. On the NVIDIA GB10 this reaches about **10-12x** at batch size 16 (CPU about 3-5x); the encoder is already compute-bound, so batching it gives no throughput win. CTC has no autoregressive decode, so batching does not apply to standalone CTC models. The batched path is bit-identical to running the clips one by one (greedy decode is deterministic). Full numbers and per-model tables are in [`benchmarks/BENCHMARK.md`](benchmarks/BENCHMARK.md#batched-decode-throughput). +The win is on the **decode** side. A transducer (TDT/RNN-T) decodes autoregressively with tiny per-step prediction-LSTM and joint GEMMs; one clip launches hundreds of these matvec-sized kernels and leaves the GPU mostly idle between launches. Decoding N clips together coalesces each step into one batched GEMM, so the device stays busy. On the NVIDIA GB10 this reaches about **10-12x** at batch size 16 (CPU about 3-5x); the encoder is already compute-bound, so batching it gives no throughput win. CTC has no autoregressive decode, so batching does not apply to standalone CTC models. The batched path is not guaranteed to be bit-identical to running the clips one by one: ggml uses a different matmul kernel for batch sizes above 1, so logits agree to about 1e-4 rather than exactly. The emitted token sequences were identical on the test clips, and the tests compare with a tolerance. Full numbers and per-model tables are in [`benchmarks/BENCHMARK.md`](benchmarks/BENCHMARK.md#batched-decode-throughput). Measure it yourself: diff --git a/docs/conversion.md b/docs/conversion.md index 742185d..7e7594c 100644 --- a/docs/conversion.md +++ b/docs/conversion.md @@ -160,6 +160,36 @@ reversed relative to numpy, so the GGUF tensor shape reads `[257, 80, 1]`. The Hann `window` buffer (`preprocessor.featurizer.window`, shape `(win_length,)`) is exported the same way. +## HF safetensors checkpoints (`convert_hf_parakeet_to_gguf.py`) + +`scripts/convert_hf_parakeet_to_gguf.py` imports checkpoints in the transformers +`ParakeetForTDT` layout that derive from `nvidia/parakeet-tdt-0.6b-v3` +(`moondream/parakeet-ultra`, `moondream/parakeet-redux`). It needs `numpy` and +`gguf` but not NeMo. The HF repos carry no mel filterbank, window or +SentencePiece vocab, so those and all KV metadata come from a v3 GGUF made by +`convert_parakeet_to_gguf.py`. The HF `config.json` is checked against that +template and every tensor shape is checked too; a mismatch aborts the run. +HF tensor names are mapped back to the verbatim NeMo names described below. + +| Argument | Meaning | +|---|---| +| `--hf DIR` | local directory with `model.safetensors`, `config.json` (and `ternary.json` for Redux) | +| `--template GGUF` | v3 GGUF that supplies KVs, vocab, filterbank and window | +| `--output GGUF` | output path | +| `--dtype f32\|f16\|q8_0` | storage type of ordinary linear weights (default `f32`) | +| `--name NAME` | `general.name` (default: the `--hf` directory name) | +| `--ternary dequant\|keep` | `dequant` (default) expands ternary weights to ordinary ones. `keep` stores them packed, see `docs/ternary.md` | +| `--vad keep\|drop` | keep (default) or drop the `vad_head.*` tensors and `parakeet.vad.*` KVs when the checkpoint has them | + +Extra GGUF content written by this converter: + +| Key or tensor | Written when | Meaning | +|---|---|---| +| `parakeet.ternary.present` (bool), `parakeet.ternary.group_size` (u32, 128) | `--ternary keep` | the file holds packed ternary linears | +| `.qweight` (I8), `.scales` (F16) | `--ternary keep` | replace `.weight` of each ternary linear | +| `parakeet.vad.present` (bool), `parakeet.vad.d_in`, `parakeet.vad.hidden`, `parakeet.vad.kernel` (u32), `parakeet.vad.frame_sec` (f32) | the checkpoint has a VAD head and `--vad keep` | shape of the head; the loader reads them into `ParakeetConfig::vad` | +| `vad_head.proj.*`, `vad_head.ctx.*`, `vad_head.out.*` (weight and bias) | same | the VAD head tensors, written as stored | + ## Worked example — `parakeet-tdt_ctc-110m` ``` diff --git a/docs/parity.md b/docs/parity.md index 2c00643..c284191 100644 --- a/docs/parity.md +++ b/docs/parity.md @@ -389,6 +389,63 @@ heads — and the C++ port reproduces each head exactly, including the second he # -> MODEL nvidia/parakeet-tdt_ctc-1.1b HEAD rnnt arch=hybrid_tdt_ctc xscaling=false WER 0.0000 ... PASS ``` +## Ultra and Redux (moondream, HF safetensors) + +`moondream/parakeet-ultra` (F16) and `moondream/parakeet-redux` (ternary +encoder) share the v3 architecture and are converted with +`scripts/convert_hf_parakeet_to_gguf.py` (see `docs/conversion.md`). There is no +NeMo baseline for these checkpoints, so parity means the transcript on +`tests/fixtures/speech.wav` equals the reference transcript in `AGENTS.md`, +and the packed ternary kernels agree with the dequantized model. + +| Model | GGUF form | Kernel | Transcript equals the reference transcript in AGENTS.md | +|---|---|---|---| +| parakeet-ultra | F16 | n/a | yes | +| parakeet-redux | packed ternary (`--ternary keep`) | scalar | yes | +| parakeet-redux | packed ternary (`--ternary keep`) | avx2 | yes | +| parakeet-redux | packed ternary (`--ternary keep`) | vnni | yes | +| parakeet-redux | dequantized F16 | n/a | yes | + +WER on long audio. Measured on three synthetic long-form clips per set, built from +LibriSpeech utterances of `benchmarks/librispeech_manifest.tsv` (30 utterances, +218 to 354 s each, gaps of low-level noise), not on TED-LIUM or other real long +recordings. Reference = the joined manifest texts; WER from `scripts/asr_metrics.py` +(case and punctuation normalized); `parakeet-cli transcribe --decoder tdt`, plain +single pass and `--vad`. Mean over three clips, percent. Details, per-clip values and the +parameter sweep are in `docs/ternary.md`. + +| Model | Gap between utterances | Plain WER | `--vad` WER | +|---|---|---:|---:| +| parakeet-ultra F16 | 0.45 s | 1.71 | 1.69 | +| parakeet-ultra F16 | 0.16 s | 1.69 | 1.78 | +| parakeet-ultra F16 | none | 1.63 | 1.84 | +| parakeet-redux packed ternary | 0.45 s | 1.97 | 1.92 | +| parakeet-redux packed ternary | 0.16 s | 1.95 | 1.69 | +| parakeet-redux packed ternary | none | 1.83 | 1.71 | + +Ultra Q8_0 and dequantized Redux were not measured on the long-form sets. On the 100 LibriSpeech +utterances (`docs/ternary.md`) Ultra Q8_0 has 1.71 percent and Redux packed 1.96 percent. + +Speed numbers are in `docs/ternary.md`. + +Multilingual and real long-form results (measured, no NeMo baseline; details, commands and caveats in +`docs/ternary.md`, section "Multilingual and long-form validation"). Scoring uses the plain `normalize` from +`scripts/asr_metrics.py`, not the Open ASR Leaderboard normalizer, so compare models with each other and not with +upstream cards. + +| Check | v3 F16 | Ultra F16 | Redux packed | Redux dequantized F16 | +|---|---:|---:|---:|---:| +| FLEURS, 25 languages x first 50 test utterances, mean WER percent | 12.77 | 10.71 | 12.09 | 11.99 | +| TED-LIUM long-form, 11 talks, plain single pass, mean WER percent | 4.40 | 3.73 | 4.32 | 5.07 | +| TED-LIUM long-form, `--vad`, mean WER percent | no VAD head | 3.64 | 4.38 | 5.17 | + +The 5.5 s clip among the 11 talks skews the Redux dequantized means; without it the packed and dequantized `--vad` +means are 4.82 and 4.85. Without the short clip the plain means are 4.75 packed and 4.74 dequantized. Ultra against the transformers `ParakeetForTDT` reference (fp32, CPU) on 60 FLEURS utterances +(en_us, de_de, fr_fr): 56 of 60 transcripts identical after normalization, WER of ours against HF 0.33 percent. +Packed Redux single-pass used to segfault above 8192 encoder frames (about 11 minutes, local attention paths without the +packed branch); fixed, and a 714 s clip gives the same transcript as the dequantized GGUF. On the seven talks that crashed, packed plain +WER is 4.92 against 4.91 dequantized, and the packed transcripts differ from the dequantized ones by 0.28 percent of words. + ## Test suite status `ctest --test-dir build --output-on-failure` (with `PARAKEET_TEST_GGUF`, diff --git a/docs/quantization.md b/docs/quantization.md index 558017e..fbd8ea8 100644 --- a/docs/quantization.md +++ b/docs/quantization.md @@ -133,6 +133,12 @@ remain in place as a safety net — those tensors are never quantized. --- +### Ternary tensors + +Packed ternary GGUFs (`--ternary keep`, see `docs/ternary.md`) store their +linears as `.qweight` (I8) and `.scales` (F16). `parakeet-cli quantize` never +re-quantizes them; it copies both tensors verbatim. + ## Measured size + WER WER is word-level vs NeMo (`scripts/validate_vs_nemo.py` on diff --git a/docs/ternary.md b/docs/ternary.md new file mode 100644 index 0000000..7ea233d --- /dev/null +++ b/docs/ternary.md @@ -0,0 +1,575 @@ +# Ternary encoders (moondream/parakeet-redux) + +## What Redux is + +`moondream/parakeet-redux` is a Parakeet TDT 0.6B v3 derivative whose encoder +linear layers are ternary: every weight is -1, 0 or +1 times a per-group scale. +Its sibling `moondream/parakeet-ultra` has the same shapes but ordinary F16 +weights, so it needs no special support and loads like any other v3 GGUF. + +In the encoder the ternary layers are the two FFN pairs (`linear1`, `linear2` +in each of `feed_forward1` and `feed_forward2`), the attention projections +(`linear_q/k/v/out/pos`) and the two pointwise convolutions, which is 11 +layers per block and 264 in the 24 block encoder. Norms, biases, the +depthwise conv, the subsampling stack, the prediction net and the joint stay +in ordinary float tensors. + +## GGUF form + +Produce it with `scripts/convert_hf_parakeet_to_gguf.py --ternary keep` (see +`docs/conversion.md`). For each ternary linear `` the file holds: + +| Tensor | Type | Content | +|---|---|---| +| `.qweight` | I8 | the upstream bytes, N rows of ceil(K/5) bytes (5 base-3 digits per byte) | +| `.scales` | F16 | one scale per row and 128 input columns, shape N x K/128 | + +and no `.weight`. Two KVs mark the file: `parakeet.ternary.present` +(bool) and `parakeet.ternary.group_size` (u32, 128). The loader reads them into +`ParakeetConfig::ternary`. A packed Redux GGUF is 213.3 MB on disk, against +1441.0 MB for the v3 F16 file. + +## Runtime path + +1. At load time `ternary_prepare` repacks every `qweight` once into a + kernel layout with 2 bits per weight (the same size as the upstream + bytes). Rows are grouped in blocks of 16, padded with zeros at the end. + For each block and each run of 16 columns there are 64 bytes: byte + 4*i + j holds the codes of row i at columns 4p + j of the run, for + p = 0..3, in its four bit pairs. The kernels keep one output row per + vector lane, so a group sum for 16 rows comes out in one vector with no + horizontal reduction. The exact layout is in `TernaryWeight` in + `src/ternary.hpp`. Nothing is dequantized per call. + The repack decodes the upstream bytes with a 256 entry table (five + trits per byte) and weights are repacked in parallel on up to + min(cores, 8) threads. `PARAKEET_REPACK_THREADS=N` overrides the count; 1 + forces the serial path. The output is byte-identical to the earlier scalar + code, which `tests/test_ternary.cpp` keeps as `repack_reference`. Measured + on packed Redux (`parakeet-cli bench`, `load_ms`, `taskset -c 0-7`, 5 runs, + load average 2.3 to 3.0): 866 ms min / 868 ms median before, 74 ms min / + 76 ms median after; with `PARAKEET_REPACK_THREADS=1`, 204 ms min / 208 ms + median. In-process, 264 weights of 4096x1024 and 1024x4096 take 1414 ms + (old code) against 210 ms (table, one thread) and 29 ms (8 threads); + `bench_ternary repack` reproduces this. +2. At run time the activations of each ternary linear are quantized to int8 + per token (one float scale and one int32 sum per 128-column group). The + scalar version defines the bytes; the AVX-512 and AVX2 versions write the + same bytes. +3. The layer is a pair of ggml custom ops: one quantizes the activations, one + runs the integer matmul and applies the scales. The dot product uses + `code * int8` sums with the constant offset removed through the group sums. +4. The kernel is picked once per process: AVX-512 VNNI, then NEON dot product, + then AVX2, then a scalar reference. All kernels are bit-identical to the + scalar reference. `PARAKEET_TERNARY_KERNEL=scalar|avx2|vnni|neon` forces one + (an unavailable name falls back to automatic selection with a log line). + The activation quantizer is picked by CPU features only. + +## Limits + +- A packed tensor anywhere in the file (any name ending in `.qweight`) while + `parakeet.ternary.present` is false is refused at load, and so is the flag set with no packed + tensors. Every tensor name is scanned, not only layer 0. +- CPU only. Loading a packed GGUF with a GPU backend active fails with a + message; re-convert with `--ternary dequant` to use a GPU. +- The kernel is chosen at run time. x86-64 with AVX2 or AVX-512 VNNI and aarch64 with the + dot-product extension (dotprod) get SIMD kernels. MSVC builds, Windows on ARM and aarch64 without + dotprod select the scalar kernel (about 1 GMAC/s); the load logs a warning when that happens. + Re-convert with `--ternary dequant` on such machines. +- The loader keeps the original packed tensors resident next to the repacked planes, so a packed + model uses more weight memory than its file size (roughly twice for Redux, not measured). +- No cache-aware streaming. `StreamingEncoder` rejects packed GGUFs; use a + `--ternary dequant` file for streaming. +- `parakeet-cli quantize` never touches ternary tensors. Running it on a packed + file copies `.qweight` (I8) and `.scales` (F16) verbatim. + +## Measured speed + +### How the numbers were measured + +The numbers come from this repository's standard build. Its CMake configure +applies the in-tree ggml patches from `third_party/ggml-patches` +(`scripts/apply_ggml_patches.sh`); the ggml submodule pin itself is unchanged. +Results on an unpatched ggml were not measured. + +Machine: AMD Ryzen 9 9950X3D (16 cores, Zen 5, AVX-512 VNNI), CPU backend, +2026-09-29. The end-to-end numbers were taken at commit f228376, the microbench +at 8b66c63 (which only changes the compile target of the activation +quantizer). Ultra has the same architecture and shapes as v3, so `ultra-q8_0` +(converted from the HF weights with `--dtype q8_0`) stands in for a Q8_0 v3. +The dequantized Redux was also converted to Q8_0 the same way. Every timed +command was serialized with `flock /tmp/pk-bench.lock`. + +### Per utterance, LibriSpeech + +``` +taskset -c 0-7 build/examples/cli/parakeet-cli bench --model \ + --manifest benchmarks/librispeech_manifest.tsv --decoder tdt --threads 8 --json +``` + +100 utterances, 901.1 s of audio, one utterance at a time. RTF is total audio +over the summed per-utterance `proc_ms` of the JSON, median of 3 full passes. +WER is against the manifest text after lowercasing and replacing every +character other than a-z and the apostrophe with a space; it is from a quick +script, not `validate_vs_nemo.py`, and is identical across passes. The kernel +rows set `PARAKEET_TERNARY_KERNEL`; "automatic" leaves it unset, which picks +vnni here. + +| Model | Form | RTF pass 1 / 2 / 3 | Median RTF | WER | +|---|---|---|---:|---:| +| parakeet-redux | packed ternary, automatic (vnni) | 74.8 / 75.6 / 78.0 | 75.6 | 1.96% | +| parakeet-redux | packed ternary, avx2 | 59.6 / 59.0 / 58.8 | 59.0 | 1.96% | +| parakeet-redux | dequantized F16 | 46.6 / 46.1 / 45.6 | 46.1 | 1.92% | +| parakeet-redux | dequantized Q8_0 | 42.7 / 42.3 / 42.7 | 42.7 | 1.96% | +| parakeet-ultra | Q8_0 | 42.5 / 43.1 / 44.2 | 43.1 | 1.71% | +| parakeet-redux | packed ternary, previous kernel (b4164da) | 40.3 / 40.2 / 40.1 | 40.2 | 1.96% | + +The last row is the kernel before the current one (one row per call, 256-bit +vectors, scalar activation quantization), built from commit b4164da and run in +the same session. The transcripts of the old and the new kernel are identical +for all 100 utterances, as expected from bit-identical kernels. + +### Long clip, one utterance + +`transcribe --decoder tdt` on `benchmarks/audio/diverse/i_have_a_dream.wav` +(180 s), whole process wall time including model load, median of 5 runs: + +``` +taskset -c 0-7 build/examples/cli/parakeet-cli transcribe --model \ + --input benchmarks/audio/diverse/i_have_a_dream.wav --decoder tdt --threads 8 +``` + +RTF is 180 s divided by the median. + +| Model | Form | Size on disk | 8 threads (cores 0-7) | RTF | 16 threads (cores 0-15) | +|---|---|---:|---:|---:|---:| +| parakeet-redux | packed ternary, vnni kernel | 213.3 MB | 7.93 s | 22.7 | 7.45 s | +| parakeet-redux | packed ternary, avx2 kernel | 213.3 MB | 8.52 s | 21.1 | 8.01 s | +| parakeet-redux | dequantized F16 | 1441.9 MB | 8.73 s | 20.6 | 7.75 s | +| parakeet-redux | dequantized Q8_0 | 941.5 MB | 9.04 s | 19.9 | 8.11 s | +| parakeet-ultra | Q8_0 | 941.5 MB | 9.04 s | 19.9 | 8.29 s | +| parakeet-ultra | F16 | 1441.9 MB | 8.63 s | 20.9 | 7.83 s | +| parakeet-tdt-0.6b-v3 | F16 | 1441.0 MB | 32.79 s | 5.5 | 29.29 s | + +The v3 F16 row read its model from a network share, and its wall time includes +that model load, so it is not comparable with the other rows. + +The scalar kernel is a reference only (116.39 s for this clip in one run at +commit c8499e6, with an earlier and faster form of the reference). + +### Single-thread kernel throughput + +``` +taskset -c 2 build/tests/bench_ternary 4096 1024 200 10 +taskset -c 2 build/tests/bench_ternary 1024 4096 200 10 +taskset -c 2 build/tests/bench_ternary 1024 4096 1000 10 +``` + +The ggml rows are `ggml_mul_mat` with an F32 activation matrix on one thread, +so they include ggml's own activation quantization; the ternary rows time only +the matmul kernel. `bench_ternary` times `ternary_quant_rows` on its own line. + +| N x K, T | scalar | ternary avx2 | ternary vnni | ggml Q8_0 | ggml F16 | ternary quant | +|---|---:|---:|---:|---:|---:|---:| +| 4096 x 1024, T=200 | 1.04 | 184.12 | 546.86 | 87.51 | 121.21 | 0.014 ms | +| 1024 x 4096, T=200 | 1.05 | 184.38 | 553.21 | 91.28 | 129.36 | 0.060 ms | +| 1024 x 4096, T=1000 | 1.05 | 183.27 | 492.23 | 89.85 | 125.12 | 0.408 ms | + +Kernel columns are GMAC/s. For comparison, the previous kernel (b4164da) +measured 81.91, 87.01 and 82.57 GMAC/s (vnni) and 78.14, 78.63 and 66.79 +(avx2) with the same commands. The quantization of the activations took 1.405 ms +for K=4096, T=200 before it was vectorized, about as long as the new matmul. +Repeated runs of the same commands moved the vnni numbers by up to about 10 +percent (the T=1000 shape measured between 492 and 542 over five runs). + +### What the data say + +- The vnni kernel does 490 to 550 GMAC/s on one core, 6.0 to 6.7 times the + previous kernel and 3.9 to 6.3 times ggml's Q8_0 and F16 mul_mat (546.86/87.51 = 6.25 at the top, 492.23/125.12 = 3.93 at the bottom). At a + clock of about 5 GHz (not measured) that is about 110 int8 multiply-adds + per cycle, close to two 512-bit `vpdpbusd` per cycle. The gain comes from + keeping one output row per vector lane (no horizontal reduction per group), + 512-bit vectors, 12 independent int32 accumulators, and unpacking the 2-bit + weights once per 4 activation rows. +- Per utterance on LibriSpeech, packed ternary is now the fastest form: median + RTF 75.6 against 46.1 for the same model in F16 and 43.1 for Ultra Q8_0, + 1.6 to 1.8 times faster. The three passes of each row are within 5 percent + of each other, well inside that margin. WER is unchanged at 1.96 percent. +- On the 180 s clip the gain is small: 7.93 s against 8.73 s (F16) and + 9.04 s (Q8_0), 9 to 12 percent less time. The linears are a smaller share of this run, + which includes model load and attention over a long sequence. +- The avx2 kernel (184 GMAC/s, 59.0 RTF) is also faster than the ggml F16 and + Q8_0 paths. It accumulates a whole group in int16, which is exact here. +- Moondream reports 113x for its own Photon runtime with these ternary weights + on 8 Zen 5 cores, against 45x for parakeet.cpp Q8_0. We did not run Photon; + our 75.6 is not measured under the same conditions and is not a comparison + with that number. +- The NEON kernel uses the same layout and is bit-identical to the scalar + reference under `qemu-aarch64`; its speed on real ARM hardware has not been + measured. +- The v3 F16 row is 3.8 times slower than Ultra and Redux F16 with identical + shapes on the long clip. That gap was not investigated. + +Transcripts on `tests/fixtures/speech.wav` are identical across the scalar, +avx2 and vnni kernels and the dequantized F16 and Q8_0 Redux, and equal the +reference transcript in `AGENTS.md`. + +## VAD head wiring + +Ultra and Redux carry a small voice-activity head (`vad_head.proj`, `vad_head.ctx`, `vad_head.out`) on the +subsampler output. The tensor shapes fix the layer order (1x1 conv 1024 to 128, conv k=5 128 to 128, 1x1 conv 128 to 1, +sigmoid), but Moondream does not document the activations or whether `ctx` is residual. We inferred the wiring from +evidence, and `VadVariant` (`src/vad_head.hpp`) keeps the three choices switchable (`parakeet-cli vad-probe --variant N`, +bit 0 ReLU after proj, bit 1 residual, bit 2 ReLU after ctx). + +How it was probed. Test data: three clips of 12 LibriSpeech utterances each (the first 36 of +`benchmarks/librispeech_manifest.tsv`, joined with no inserted silence), `jfk.wav`, and the first 120 s of +`i_have_a_dream.wav`. Labels came from the model's own TDT word timestamps: a speech frame (80 ms) has its center inside a +word span widened by 0.04 s; a pause frame lies inside an inter-word, leading or trailing gap of at least 0.4 s, minus a +0.16 s margin at each end. Counts: 4289 speech frames, 479 pause frames, 218 ignored, identical for both models. (An +earlier probe with digital silence and white noise clips was discarded: zero padding wrecks the per-feature mel +normalization of the speech part.) Features were dumped once per clip from the F16 Ultra and the dequantized F16 Redux. + +Grid, 108 variants: tap (subsampler output, final encoder output) x input treatment (none, per-frame LayerNorm without +affine, per-frame L2 normalize times sqrt(1024)) x activation after proj (none, ReLU, SiLU) x activation after ctx (none, +ReLU, SiLU) x ctx residual (no, yes). Each is scored by the median p on speech frames, the median p on pause frames and +the ROC AUC of p separating the two. + +Top rows (speech median, pause median, AUC). All are the subsampler tap with no input treatment: + +| proj act, ctx act, residual | Ultra | Redux | mean AUC | +| --- | --- | --- | --- | +| ReLU, ReLU, no (chosen, variant 5) | 0.999, 0.437, 0.930 | 1.000, 0.197, 0.948 | 0.939 | +| ReLU, SiLU, no | 0.999, 0.437, 0.930 | 1.000, 0.197, 0.948 | 0.939 | +| SiLU, ReLU, no | 0.999, 0.437, 0.930 | 1.000, 0.197, 0.948 | 0.939 | +| SiLU, SiLU, no | 0.999, 0.437, 0.930 | 1.000, 0.197, 0.948 | 0.939 | +| ReLU, ReLU, yes | pause median 1.000 (collapses) | 1.000, 0.200, 0.959 | | + +The SiLU rows tie with ReLU because the pre-activations are huge and saturated. No variant met the strict rule (AUC of at +least 0.90, speech median at least 0.8, pause median at most 0.2, on both models); variant 5 misses it only on the Ultra +pause median (0.437). + +Refuted. The final encoder output as the tap scores AUC 0.17 to 0.67 for every treatment and activation, so it is at +chance or inverted. LayerNorm and L2 input normalization do not help: the medians collapse to about 0.5 for both classes +with equal or lower AUC. + +Decision: subsampler tap, no input treatment, ReLU after proj, ReLU after ctx, no residual (variant 5, the default of +`VadVariant`). + +Caveats: +- Pause detection is weak. Pause frames with p below 0.5: 58 percent on Ultra, 94 percent on Redux (the same for gaps of + 0.8 s and 1.5 s or more). Speech frames above 0.5: 92.8 percent on Ultra, 92.9 percent on Redux. +- Ultra and Redux disagree on the residual: Ultra needs none (with it the pause median is 1.000), Redux scores best with it + (AUC 0.959 against 0.948). +- The wiring is inferred, not documented by Moondream. Something may still differ from training. +- The default threshold is 0.5. Per-model thresholds may differ (Ultra about 0.5 to 0.7, Redux about 0.3). +- Long-form validation is in the next section. The segmenter falls back to hard cuts when it finds no pause, so a weak + pause response degrades gracefully. + +## Long-form WER with and without the VAD + +`--vad` cuts long audio at pauses found by the VAD head and transcribes each segment on its own +(`segment_by_vad`, defaults: threshold 0.5, min pause 0.32 s, max segment 30 s, min segment 8 s). Without it, the model +sees the whole clip in one pass. + +How it was measured. `scripts/make_longform.py` joins the first 90 utterances of `benchmarks/librispeech_manifest.tsv` +into three clips of 30 utterances each (218 to 354 s each), with a known reference +(the joined manifest texts). Utterances are separated by Gaussian noise at about -55 dBFS (fixed seed) rather than +digital zeros, because digital silence distorts the per-feature mel normalization. Three sets: a 0.45 s gap, a 0.16 s gap +(below the 0.32 s minimum pause on purpose), and no inserted gap (only the natural utterance edges). These are synthetic +long-form clips built from LibriSpeech read speech, not TED-LIUM or other real long recordings. +`scripts/eval_vad_longform.py` runs `parakeet-cli transcribe --decoder tdt --threads 8` under `taskset -c 0-7` with and +without `--vad` and scores with `scripts/asr_metrics.py` `wer` (case and punctuation normalized). CPU pinning and the shared bench lock are applied by the caller through the script's `--prefix` option, not by the +script. The command used: + +``` +python3 scripts/eval_vad_longform.py --model --dir --glob 'longform_0p45_*.wav' \ + --prefix 'flock /tmp/pk-bench.lock taskset -c 0-7' +``` + +(`--threads 8` is the script default; sweep runs add `--skip-plain --vad-arg=--vad-threshold=0.3` and similar.) The models are the +Ultra F16 GGUF and the packed ternary Redux GGUF (`--ternary keep`, native kernel). + +WER per clip, plain single pass vs `--vad` (percent, three clips per set, then the mean): + +| Model | Gap | Clip 0 | Clip 1 | Clip 2 | Mean | +|---|---|---:|---:|---:|---:| +| Ultra F16, plain | 0.45 s | 0.47 | 2.39 | 2.29 | 1.71 | +| Ultra F16, `--vad` | 0.45 s | 0.47 | 2.21 | 2.39 | 1.69 | +| Ultra F16, plain | 0.16 s | 0.31 | 2.58 | 2.18 | 1.69 | +| Ultra F16, `--vad` | 0.16 s | 0.62 | 2.21 | 2.50 | 1.78 | +| Ultra F16, plain | none | 0.31 | 2.39 | 2.18 | 1.63 | +| Ultra F16, `--vad` | none | 0.62 | 2.39 | 2.50 | 1.84 | +| Redux packed, plain | 0.45 s | 0.62 | 2.58 | 2.72 | 1.97 | +| Redux packed, `--vad` | 0.45 s | 0.47 | 2.58 | 2.72 | 1.92 | +| Redux packed, plain | 0.16 s | 0.62 | 2.39 | 2.83 | 1.95 | +| Redux packed, `--vad` | 0.16 s | 0.31 | 2.03 | 2.72 | 1.69 | +| Redux packed, plain | none | 0.62 | 2.03 | 2.83 | 1.83 | +| Redux packed, `--vad` | none | 0.62 | 2.21 | 2.29 | 1.71 | + +Reading. The change in mean WER from plain to `--vad` runs in both directions and stays within about 0.2 points: +Ultra -0.02, +0.09 and +0.21 (0.45 s, 0.16 s, no gap); Redux -0.05, -0.26 and -0.12. The single largest loss (+0.21, Ultra, +no gap) is marginally over the 0.2-point investigation limit. The Redux gains and the Ultra losses are the same size, so +neither direction is a real effect on three clips per set: VAD segmentation does not change WER meaningfully on these +clips. One clip is 645 to 919 words, so 0.1 point is about one word. Inspecting the worst Ultra clip (no gap, clip 0, 10 segments) shows every cut lands between words, none +inside one, and the differing words are spelling variants (`tail`/`tale`, `honour`/`honor`) and rare names that flip +between the two runs. The segment boundaries there were 20.32, 48.48, 70.00, 98.48, 127.12, 156.56, 179.92, 204.16 and +229.28 s. The result is a wash, not a win: on these clips VAD does not measurably help or hurt accuracy. Its benefit is +bounded memory and time on audio of any length (the 30 s cap), not a lower WER. + +Parameter sweep on Ultra F16 (mean `--vad` WER, percent, one variable at a time from the defaults): + +| Setting | 0.45 s gap | 0.16 s gap | no gap | mean of the three | +|---|---:|---:|---:|---:| +| threshold 0.3 | 1.69 | 1.67 | 1.89 | 1.75 | +| threshold 0.5 (default) | 1.69 | 1.78 | 1.84 | 1.77 | +| threshold 0.7 | 1.69 | 1.83 | 1.89 | 1.80 | +| min pause 0.16 s | 1.69 | 1.78 | 1.84 | 1.77 | +| min pause 0.32 s (default) | 1.69 | 1.78 | 1.84 | 1.77 | +| min pause 0.64 s | 1.95 | 1.79 | 1.85 | 1.86 | + +Threshold 0.3 is the only setting ahead of the default. On Redux packed it gives 1.97, 1.65 and 1.63 (mean 1.75) +against 1.92, 1.69 and 1.71 (mean 1.77) for the default. Both gains are 0.02 points, under half a word per clip set, and +Redux gets worse on the 0.45 s gap set, so the defaults stay. A min pause of 0.64 s is worse on the 0.45 s gap set for +Ultra (1.95); probably a longer required pause makes the segmenter fall back to hard cuts more often, which was not measured. A min pause of 0.16 s +changes nothing here. + +Known limits: +- A 30 s window with no pause of 0.32 s makes a hard cut, which can land inside a word. +- Offline only: no streaming with `--vad`. +- The VAD head detects pauses weakly (see the caveats above), so many cuts are hard cuts or land at the model's best guess. +- Ternary kernels are CPU only. +- The figures in this section are from synthetic LibriSpeech clips (218 to 354 s); real talks are in the multilingual and long-form section below. + +## Multilingual and long-form validation + +Everything above is English (LibriSpeech and synthetic long clips built from it). Ultra is a 25-language model, and the +upstream cards report FLEURS and TED-LIUM long-form results, so this section measures our engine on both. Nothing was +tuned, no `src/` code changed, and every number below comes from the commands listed here. The scoring, the subsets and +the caveats matter for how to read the tables, so read those first. + +### Scoring and subset rules + +- Scoring uses `normalize` from `scripts/asr_metrics.py` on both reference and hypothesis (NFKC, lowercase, punctuation + replaced by a space, whitespace collapsed). It is the same for every model. It is NOT the Open ASR Leaderboard + normalizer (no number spelling, no text-normalizer for other languages), so the absolute numbers here are not + comparable with the upstream model cards. What matters is the direction and size of the differences between our + models on the same subset. +- Upstream's own figures, from their pipeline on the full test splits, quoted only as theirs: FLEURS average 11.62 for + v3, 9.55 for Ultra and 10.56 for Redux; TED-LIUM long-form 2.71 for v3, 1.94 for Ultra and 2.51 for Redux. +- FLEURS: the first 50 utterances of the `test` split of each language in dataset (streaming parquet) order, no + shuffling and no filtering. The reference is the dataset's `raw_transcription` (natural case and punctuation). + 25 languages, 1250 utterances, 26784 reference words, 14580.6 s of audio. Audio is resampled to 16 kHz mono int16. + Because each language has about 1000 words, one word is about 0.1 point, so differences under roughly 1 point on a + single language are noise. WER per language is a corpus WER (total edits over total reference words); the mean is the + unweighted mean over the 25 languages. +- The config ids all exist as listed: bg_bg, hr_hr, cs_cz, da_dk, nl_nl, en_us, et_ee, fi_fi, fr_fr, de_de, el_gr, + hu_hu, it_it, lv_lv, lt_lt, mt_mt, pl_pl, pt_br, ro_ro, ru_ru, sk_sk, sl_si, es_419, sv_se, uk_ua. +- TED-LIUM: `distil-whisper/tedlium-long-form`, `test` split, 11 full talks (8905 s in total, 2.5 hours). The reference + is the `text` column with tags such as `` removed. One of the 11 is a 5.5 s clip with 24 words + (`DanBarber_2010_S103`), so it moves a per-talk mean by a lot; means are given with and without it. +- Models: v3 F16 (`tdt-0.6b-v3-f16.gguf`), Ultra F16, Redux packed ternary (`--ternary keep`, native kernel) and + Redux dequantized F16 (`--ternary dequant`, isolates the int8 activation and packed kernel path). All with + `--decoder tdt`. + +### FLEURS-25, WER percent + +| FLEURS config | v3 F16 | Ultra F16 | Redux packed | Redux dequantized F16 | +|---|---:|---:|---:|---:| +| bg_bg | 11.24 | 9.23 | 10.37 | 10.98 | +| cs_cz | 11.93 | 9.61 | 10.62 | 9.71 | +| da_dk | 17.96 | 16.18 | 16.44 | 16.71 | +| de_de | 4.75 | 3.67 | 4.92 | 4.83 | +| el_gr | 36.14 | 33.01 | 32.13 | 32.29 | +| en_us | 5.30 | 4.65 | 6.23 | 6.04 | +| es_419 | 3.10 | 2.38 | 3.42 | 3.65 | +| et_ee | 17.83 | 15.21 | 12.47 | 13.34 | +| fi_fi | 13.06 | 9.77 | 11.84 | 12.21 | +| fr_fr | 4.23 | 4.45 | 10.31 | 10.09 | +| hr_hr | 11.80 | 8.90 | 9.67 | 9.19 | +| hu_hu | 16.18 | 12.69 | 18.48 | 18.08 | +| it_it | 1.88 | 2.11 | 3.54 | 3.69 | +| lt_lt | 22.03 | 18.57 | 20.30 | 19.98 | +| lv_lv | 23.56 | 19.08 | 14.39 | 14.61 | +| mt_mt | 21.20 | 16.59 | 15.46 | 15.29 | +| nl_nl | 7.94 | 7.15 | 10.86 | 10.41 | +| pl_pl | 8.71 | 7.69 | 11.54 | 11.43 | +| pt_br | 4.55 | 3.99 | 5.58 | 5.34 | +| ro_ro | 12.73 | 11.31 | 12.56 | 12.90 | +| ru_ru | 5.68 | 5.11 | 8.14 | 7.85 | +| sk_sk | 9.22 | 7.01 | 9.42 | 9.22 | +| sl_si | 24.12 | 19.00 | 20.81 | 20.17 | +| sv_se | 17.39 | 14.87 | 14.57 | 14.07 | +| uk_ua | 6.81 | 5.43 | 8.09 | 7.55 | +| Mean over 25 languages | 12.77 | 10.71 | 12.09 | 11.99 | + +Reading: + +- Ultra beats v3 on 23 of 25 languages, mean 10.71 against 12.77 (2.07 points lower on average). It loses slightly on + fr_fr (4.45 against 4.23) and it_it (2.11 against 1.88), which are within noise at this sample size. The size of the + gain is close to upstream's (11.62 to 9.55), which is a useful sign that the conversion did not lose anything, but + the two pipelines differ so this is a direction check, not a reproduction. +- Redux is better than v3 on 13 languages (mostly Baltic, Slavic, Uralic and Greek: bg, cs, da, el, et, fi, hr, lt, lv, + mt, ro, sl, sv) and worse on the rest; the mean is 12.09 against 12.77. It loses to Ultra on 20 of 25. The largest + gaps to Ultra are fr_fr (10.31 against 4.45, +5.9), hu_hu (+5.8), pl_pl (+3.9), nl_nl (+3.7), ru_ru (+3.0), uk_ua + (+2.7) and sk_sk (+2.4). French and Polish match what the upstream card says Redux gives up. Redux is ahead of Ultra + on el_gr, et_ee, lv_lv, mt_mt and sv_se, with lv_lv (14.39 against 19.08) and et_ee (12.47 against 15.21) the biggest. + I did not check why; it may be a difference in the training mix. +- Int8 activations. Packed Redux (12.09) against dequantized Redux F16 (11.99) is +0.10 points on the mean, with + per-language differences from -0.87 to +0.91 and no direction (packed is lower on 9 languages, higher on 16). With + about 1000 words per language that is noise. On this data the packed ternary path costs no measurable WER against the + dequantized weights, and the speed benefit is in the speed section above. +- el_gr is high for every model (32 to 36). The FLEURS Greek references and the model outputs likely disagree on + normalization details (accents, final sigma, number forms); I did not investigate, and it does not change the + comparison between models on the same references. + +### TED-LIUM long-form, WER percent + +Plain is one single pass over the whole talk; VAD is `--vad` with default options +(threshold 0.5, min pause 0.32 s, max segment 30 s, min segment 8 s). The `--vad` option needs the VAD head, which v3 +does not have (`model has no VAD head`), so v3 has only the plain column. The segment count was not recorded. + +| Talk (length) | v3 plain | Ultra plain | Ultra VAD | Redux packed plain | Redux packed VAD | Redux deq plain | Redux deq VAD | +|---|---:|---:|---:|---:|---:|---:|---:| +| AimeeMullins (1249 s) | 6.65 | 3.44 | 3.41 | 4.18 | 4.34 | 4.08 | 4.21 | +| BillGates (1506 s) | 7.07 | 5.79 | 5.74 | 6.18 | 6.34 | 6.32 | 6.32 | +| DanBarber (834 s) | 5.75 | 5.46 | 4.90 | 6.47 | 6.79 | 6.59 | 6.83 | +| DanBarber_2010_S103 (5.5 s) | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 8.33 | 8.33 | +| DanielKahneman (1096 s) | 3.29 | 3.23 | 3.20 | 4.11 | 3.98 | 4.11 | 4.14 | +| EricMead (459 s) | 5.61 | 4.91 | 4.84 | 5.35 | 5.16 | 5.23 | 5.16 | +| GaryFlake (345 s) | 2.90 | 3.16 | 3.08 | 3.52 | 3.60 | 3.60 | 3.69 | +| JamesCameron (982 s) | 6.34 | 5.52 | 5.42 | 5.81 | 5.62 | 5.62 | 5.55 | +| JaneMcGonigal (1168 s) | 4.21 | 3.60 | 3.68 | 4.06 | 4.09 | 4.06 | 4.16 | +| MichaelSpecter (921 s) | 3.57 | 3.25 | 3.16 | 3.63 | 3.79 | 3.63 | 3.86 | +| RobertGupta (340 s) | 3.04 | 2.71 | 2.59 | 4.17 | 4.51 | 4.17 | 4.62 | +| Mean, all 11 | 4.40 | 3.73 | 3.64 | 4.32 | 4.38 | 5.07 | 5.17 | +| Mean, without the 5.5 s clip | 4.84 | 4.11 | 4.00 | 4.75 | 4.82 | 4.74 | 4.85 | + +Reading: + +- Ultra beats v3 on 9 of the 11 talks in a single pass (one tie on the 5.5 s clip) and on the mean (3.73 against 4.40; 4.11 against 4.84 without the + short clip). It loses only on GaryFlake (3.16 against 2.90). Direction agrees with upstream (2.71 against 1.94), but our + absolute values are about twice upstream's. Part of that gap is the normalizer: the references spell numbers out while + the model writes digits, and in the Ultra plain hypotheses 218 of about 26800 tokens contain digits (none in the + references), which alone accounts for at least 0.8 points. I did not run a number-aware normalizer, so the rest of the gap is + not explained here. +- Redux packed with `--vad`: 4.38 on all 11 (4.82 without the short clip), against Ultra VAD 3.64 (4.00) and v3 plain + 4.40 (4.84). So Redux is about where v3 is on real talks and about 0.8 points behind Ultra, the same order as upstream + (2.51 against 1.94 against 2.71). +- Packed against dequantized Redux, both with `--vad`: 4.38 against 5.17 on all 11, but 4.82 against 4.85 without the + 5.5 s clip. The all-11 gap is the 24-word clip (0 against 2 errors); on the ten real talks packed and dequantized are the same. +- VAD against plain. Ultra: VAD is better on 9 of 11 talks (one tie), mean 3.64 against 3.73 (4.00 against 4.11), gains of 0.03 + to 0.56 points (the 0.56 is DanBarber), and worse on JaneMcGonigal by 0.08. Redux dequantized: VAD is slightly worse, 5.17 against 5.07 + (4.85 against 4.74). Redux packed: plain 4.32 against VAD 4.38 (4.75 against 4.82 without the short clip), so plain is slightly better, by + 0.07 points. The honest reading is that on real talks VAD segmentation is about neutral: a small gain for Ultra, a small loss for Redux, + all under about 0.15 points on the mean and inside per-talk noise. Peak RSS (from `/usr/bin/time -v`) with `--vad` is 7 to 11 GB on the talks over 800 s, about the same as + the single pass, and drops to about 4 GB only on the three shortest talks (340 to 459 s). So on these talks VAD did not buy a memory win. +- Packed against dequantized Redux, plain single pass, after the long-audio fix (the seven talks that used to crash; packed run + at `--threads 8` with `/usr/bin/time -v`): WER 4.92 packed against 4.91 dequantized on the mean of the seven (per talk within + 0.2 points, packed lower on 2, higher on 2, equal on 3). Scored with the dequantized transcript as the reference, packed + differs by 64 of 22922 words (0.28 percent; per talk 0.05 to 0.50 percent). So on real 14 to 25 minute talks the packed + path matches the dequantized one to a few words per talk, and there is no measurable cost from int8 activations or the + packed kernel. The packed plain runs all exited 0. Longest talk (BillGates, 1506 s): 82.8 s wall, 11.67 GB peak RSS, load average 20 at the + start; the other six took 64 to 85 s except JamesCameron (318 s) and MichaelSpecter (294 s), which started at load 26 to 48; + peak RSS 6.6 to 11.7 GB, similar to the F16 models. Load averages at start were 3.5 to 48, so wall times are indicative only. +- Single-pass memory and time. Every single pass over a 5 to 25 minute talk finished on the F16 models (Ultra, dequantized Redux, + v3), peak RSS 7 to 13 GB, so the O(T^2) attention above the 8192-frame local-attention threshold did not fail for any of these + talks. Wall times were recorded but are not reported: the shared machine was heavily loaded during the runs (load average 40 to + 90 from other jobs), so times vary by more than 10x between identical runs. + +#### Finding (fixed): packed Redux crashed on single-pass audio longer than about 11 minutes + +With `redux-keep.gguf` (packed ternary), plain single pass segfaulted after about 3 s on every talk longer than about +655 s (8192 encoder frames at 80 ms): AimeeMullins, BillGates, DanBarber, DanielKahneman, JamesCameron, JaneMcGonigal and +MichaelSpecter, exit code 139. The backtrace ended in `RelPosAttention::build_graph_local_chunked`. + +Cause: above `kLocalThreshold = 8192` frames the encoder switches to the local attention paths. Three of the five +attention `linear` lambdas (`build_graph_batched_local`, `build_graph_local`, `build_graph_local_chunked`) read +only `.weight`, which a packed GGUF does not have, so `ggml_mul_mat` got a null tensor. The earlier synthetic clips +(218 to 354 s) never reach these paths. All five now go through one helper (`attn_linear` in +`src/relpos_attention.cpp`), which uses the packed kernel when `.qweight` exists and throws +`missing encoder weight ` when neither form exists. + +Verified after the fix on a 714 s clip (`speech.wav` repeated, 2112 words): packed and dequantized Redux both finish and +print byte-identical transcripts. `tests/test_ternary_long.cpp` forces the same paths on short audio +(`PARAKEET_ATT_CONTEXT=64`), for a single item and for a batch. The seven talks that crashed were then re-run +with the fixed build; their packed plain numbers are in the TED-LIUM table above. + +### Parity with the transformers reference (Ultra) + +An independent check that the Ultra conversion and engine reproduce the HF implementation. `ParakeetForTDT` is not in the +installed transformers 5.3.0, so this used a source checkout that has it (version string 5.10.0.dev0) through +`PYTHONPATH`. The HF checkpoint has no preprocessor or tokenizer files, so `scripts/hf_reference_transcribe.py` builds a +`ParakeetFeatureExtractor(feature_size=128, sampling_rate=16000)` (per-feature normalization is built in) and decodes +token ids with the piece table from our Ultra GGUF. The run is fp32 on CPU with greedy TDT decoding; ours is the F16 GGUF. +Because the token table comes from our GGUF, this checks the encoder, decoder and search, not the tokenizer. + +Same first 20 test utterances of en_us, de_de and fr_fr (60 utterances, 1509 words). Ours scored against the HF +transcript as reference, after `normalize`, and after dropping the `` piece, which the HF side prints and we omit: + +| Language | Identical after normalize | WER of ours vs HF | +|---|---:|---:| +| en_us | 20 of 20 | 0.00 | +| de_de | 18 of 20 | 0.45 | +| fr_fr | 18 of 20 | 0.48 | +| total | 56 of 60 | 0.33 | + +The four differing utterances are single-word choices near a tie (for example `Hirnschadens` against `Höhenschadens`, +`Laka` against `Lakas`, `vient` against `viant`, and `dix-sept` against `17`). The F16 weights in the GGUF against the fp32 weights in HF are a plausible cause, but I did not test that, for example by converting an +F32 GGUF. + +Both implementations scored against the FLEURS references (same scorer as above, `` left in the HF text, where the +`unk` word counts as an error): + +| Language | Ours (F16 GGUF) | HF transformers (fp32) | +|---|---:|---:| +| en_us | 3.64 | 3.64 | +| de_de | 3.56 | 3.78 | +| fr_fr | 5.13 | 6.09 | +| mean of the three | 4.11 | 4.50 | + +The fr_fr gap on the HF side comes from `` pieces that the HF text prints (four in the 20 utterances) and we drop; it is a +decoding detail, not a model difference. + +### Commands + +``` +# subsets (stream from the Hub, no token; only 16 kHz wavs, references and manifests are written, under 1 GB in total) +python3 scripts/fetch_fleurs_subset.py --out /tmp/val/fleurs --n 50 +python3 scripts/fetch_tedlium_longform.py --out /tmp/val/ted + +# FLEURS WER for the four models (bench --json keeps the hypotheses per language) +python3 scripts/eval_manifest_wer.py --out /tmp/val/fleurs_res --threads 8 \ + --model v3= --model ultra= \ + --model redux= --model reduxdeq= /tmp/val/fleurs/*/manifest.tsv + +# TED-LIUM plain and --vad, one command per model (records WER, wall time, peak RSS; a crash is recorded as FAIL) +python3 scripts/eval_longform_talks.py --model --dir /tmp/val/ted --save /tmp/val/ted_res/ + +# HF transformers reference for Ultra, then compare +PYTHONPATH=/src python3 scripts/hf_reference_transcribe.py \ + --hf-dir --gguf --out /tmp/val/spike_hf /tmp/val/spike/{en_us,de_de,fr_fr}/manifest.tsv +python3 scripts/compare_hyps.py /tmp/val/spike_res/ultra /tmp/val/spike_hf +``` + +The FLEURS run used `--threads 2` for most of the sweep instead of 8, because the machine was heavily loaded and 8 spinning +threads were much slower than 2 there; greedy decoding does not depend on the thread count except for float summation +order, which was not checked for any change in a transcript. + +Limits of this section: +- 50 utterances per language is a sample, not the full FLEURS test split; per-language numbers carry about a 1 point noise band. +- The normalizer is a plain one, so numbers, hyphenation and non-Latin scripts add errors that the leaderboard normalizer + would remove; compare models against each other, not against upstream. +- The TED-LIUM set has 11 talks, one of them 5.5 s long; a per-talk difference under 0.2 points is not meaningful. +- Timing is not reported (loaded machine). The speed numbers are in the earlier speed section. + +## Tests + +``` +ctest --test-dir build -R ternary --output-on-failure # test_ternary (kernels vs scalar), no model needed + +PARAKEET_TEST_GGUF_REDUX_KEEP= \ +PARAKEET_TEST_GGUF_REDUX_DEQ= \ + ctest --test-dir build -R test_ternary_model --output-on-failure + +build/tests/bench_ternary [N K T reps] # single thread: each kernel, the quantizer, ggml Q8_0 and F16 mul_mat +``` diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index f2fa537..7f350f6 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -1,6 +1,7 @@ #include "parakeet.h" #include "parakeet_capi.h" #include "model.hpp" +#include #include "model_loader.hpp" #include "audio_io.hpp" #include "streaming.hpp" @@ -22,6 +23,7 @@ #include "ced_tagger.hpp" #include "scene_stream.hpp" #include "scene_render.hpp" +#include "vad_head.hpp" #include #include #include @@ -195,6 +197,36 @@ static int cmd_transcribe_stream(const std::string& model, const std::string& in return 0; } +static int cmd_transcribe_vad(const std::string& model, const std::string& input, pk::Decoder dec, + const std::string& lang, bool timestamps, bool json, + const pk::SegmenterOpts& opts) { + pk::Audio audio; + if (!load_audio_arg_16k_mono(input, audio)) { + std::fprintf(stderr, "parakeet-cli: failed to load audio %s\n", input.c_str()); + 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; } + if (json || timestamps) { + pk::Transcription tr = + m->transcribe_pcm_vad_with_timestamps(audio.samples, audio.sample_rate, dec, lang, opts); + if (json) { + std::printf("%s\n", pk::transcription_to_json(tr, model_frame_sec(*m)).c_str()); + } else { + for (const pk::Word& w : tr.words) + std::printf("%.2f-%.2f %s (%.2f)\n", w.start, w.end, w.text.c_str(), w.conf); + } + } else { + std::printf("%s\n", m->transcribe_pcm_vad(audio.samples, audio.sample_rate, dec, lang, opts).c_str()); + } + } catch (const std::exception& e) { + std::fprintf(stderr, "parakeet-cli: transcribe failed: %s\n", e.what()); + return 1; + } + return 0; +} + // parakeet-cli transcribe --model --input [--decoder ctc|tdt] // [--stream] // Prints the transcript. Default decoder is chosen by arch (TDT for transducer @@ -205,6 +237,14 @@ static int cmd_transcribe(int argc, char** argv) { bool stream = false; bool timestamps = false; bool json = false; + bool vad = false; + pk::SegmenterOpts vad_opts; + double d = 0.0; + auto parse_pos = [](const char* str, double& out) { + char* end = nullptr; + out = std::strtod(str, &end); + return end != str && *end == '\0' && std::isfinite(out) && out > 0.0 && out <= 1e6; + }; bool score_norm = true; int beam_size = 0; int nbest = 0; @@ -232,6 +272,17 @@ static int cmd_transcribe(int argc, char** argv) { nbest = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--no-score-norm") == 0) { score_norm = false; + } else if (std::strcmp(argv[i], "--vad") == 0) { + vad = true; + } else if (std::strcmp(argv[i], "--vad-threshold") == 0 && i + 1 < argc) { + if (!parse_pos(argv[++i], d) || d > 1.0) { std::fprintf(stderr, "parakeet-cli: --vad-threshold must be in (0,1]\n"); return 2; } + vad_opts.threshold = (float)d; + } else if (std::strcmp(argv[i], "--vad-min-pause") == 0 && i + 1 < argc) { + if (!parse_pos(argv[++i], d)) { std::fprintf(stderr, "parakeet-cli: --vad-min-pause must be > 0\n"); return 2; } + vad_opts.min_pause_sec = d; + } else if (std::strcmp(argv[i], "--vad-max-seg") == 0 && i + 1 < argc) { + if (!parse_pos(argv[++i], d)) { std::fprintf(stderr, "parakeet-cli: --vad-max-seg must be > 0\n"); return 2; } + vad_opts.max_seg_sec = d; } } if (model.empty() || input.empty()) { @@ -239,6 +290,7 @@ static int cmd_transcribe(int argc, char** argv) { "usage: parakeet-cli transcribe --model --input " "[--decoder ctc|tdt] [--lang ] [--stream] [--timestamps] " "[--threads N] [--json] " + "[--vad [--vad-threshold F] [--vad-min-pause SEC] [--vad-max-seg SEC]] " "[--beam-size N [--nbest N] [--no-score-norm]]\n"); return 2; } @@ -247,6 +299,10 @@ static int cmd_transcribe(int argc, char** argv) { if (threads > 0) pk::set_num_threads(threads); if (stream) { + if (vad) { + std::fprintf(stderr, "parakeet-cli: --vad is offline only\n"); + return 2; + } if (beam_size != 0 || nbest != 0) { std::fprintf(stderr, "parakeet-cli: --beam-size/--nbest are offline TDT only\n"); @@ -281,6 +337,13 @@ static int cmd_transcribe(int argc, char** argv) { } } + if (vad) { + if (beam_size != 0 || nbest != 0) { + std::fprintf(stderr, "parakeet-cli: --vad works with greedy decoding only\n"); + return 2; + } + return cmd_transcribe_vad(model, input, dec, lang, timestamps, json, vad_opts); + } if (nbest != 0 && beam_size == 0) { std::fprintf(stderr, "parakeet-cli: --nbest requires --beam-size\n"); @@ -1506,6 +1569,39 @@ static int run_and_shutdown(int (*fn)(int, char**), int argc, char** argv) { return rc; } +// parakeet-cli vad-probe --model --input [--variant N] +// Prints "t_sec,p" for every 80 ms frame. For inspecting the VAD head. +static int cmd_vad_probe(int argc, char** argv) { + std::string model, input; + int variant = -1; + for (int i = 0; i < argc; ++i) { + if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) model = argv[++i]; + else if (std::strcmp(argv[i], "--input") == 0 && i + 1 < argc) input = argv[++i]; + else if (std::strcmp(argv[i], "--variant") == 0 && i + 1 < argc) variant = std::atoi(argv[++i]); + } + if (model.empty() || input.empty()) { + std::fprintf(stderr, "usage: parakeet-cli vad-probe --model --input [--variant N]\n"); + return 2; + } + pk::Audio audio; + if (!load_audio_arg_16k_mono(input, audio)) { + std::fprintf(stderr, "parakeet-cli: failed to load audio %s\n", input.c_str()); + 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; } + const pk::VadVariant v = pk::VadVariant::from_index(variant < 0 ? 0 : variant); + const std::vector p = m->vad_probabilities(audio.samples, variant < 0 ? nullptr : &v); + const float fs = m->config().vad.frame_sec; + for (size_t i = 0; i < p.size(); ++i) std::printf("%.2f,%.4f\n", (double)i * fs, p[i]); + } catch (const std::exception& e) { + std::fprintf(stderr, "parakeet-cli: vad-probe failed: %s\n", e.what()); + return 1; + } + return 0; +} + int main(int argc, char** argv) { if (argc == 2 && (std::strcmp(argv[1], "--version") == 0 || std::strcmp(argv[1], "-V") == 0)) { @@ -1526,12 +1622,16 @@ int main(int argc, char** argv) { return run_and_shutdown(cmd_bench, argc - 2, argv + 2); if (argc >= 2 && std::strcmp(argv[1], "scene") == 0) return run_and_shutdown(cmd_scene, argc - 2, argv + 2); + if (argc >= 2 && std::strcmp(argv[1], "vad-probe") == 0) + return run_and_shutdown(cmd_vad_probe, argc - 2, argv + 2); std::fprintf(stderr, "usage:\n" + " parakeet-cli vad-probe --model --input [--variant N]\n" " parakeet-cli info \n" " parakeet-cli transcribe --model --input " "[--decoder ctc|tdt] [--lang ] [--stream] [--timestamps] " "[--threads N] [--json] " + "[--vad [--vad-threshold F] [--vad-min-pause SEC] [--vad-max-seg SEC]] " "[--beam-size N [--nbest N] [--no-score-norm]]\n" " parakeet-cli quantize " "\n" diff --git a/include/parakeet_capi.h b/include/parakeet_capi.h index 05a27a3..715b0ea 100644 --- a/include/parakeet_capi.h +++ b/include/parakeet_capi.h @@ -150,6 +150,16 @@ int parakeet_capi_transcribe_pcm_batch_lang(parakeet_ctx* ctx, char* parakeet_capi_transcribe_path_json(parakeet_ctx* ctx, const char* wav_path, int decoder); +// Like parakeet_capi_transcribe_path_json, but long audio is cut at pauses found +// by the model's own VAD head into segments of at most 30 s, and the segments are +// transcribed one by one (word/token times are relative to the whole file). Audio +// of 30 s or less gives the same document as the plain function. Returns NULL and +// sets the context's last error to "model has no VAD head" when the model has no +// VAD head. It always uses the default segmenter options (30 s cap, threshold +// 0.5). Additive; no ABI bump. +char* parakeet_capi_transcribe_path_json_vad(parakeet_ctx* ctx, const char* wav_path, + int decoder); + // Batched transcription with timestamps, returning ONE malloc'd JSON string that // is a JSON ARRAY of n_clips objects, each identical in shape to // parakeet_capi_transcribe_path_json's document ({"text","words","tokens"}). diff --git a/models/MANIFEST.md b/models/MANIFEST.md index 9963c5b..e1c910e 100644 --- a/models/MANIFEST.md +++ b/models/MANIFEST.md @@ -100,6 +100,29 @@ WER (word error rate) is measured against the NeMo reference on | Q8_0 | `mudler/parakeet.cpp-parakeet-rnnt-1.1b-q8_0` | ~400 MB | **0.0** | PASS | | Q4_K | `mudler/parakeet.cpp-parakeet-rnnt-1.1b-q4_k` | ~300 MB | not yet measured | — | +### `moondream/parakeet-ultra` and `moondream/parakeet-redux` (HF safetensors, v3 shape) + +Converted with `scripts/convert_hf_parakeet_to_gguf.py --template `. +Neither is published yet; sizes are from local conversions. Transcript on +`tests/fixtures/speech.wav` matches the reference. WER on LibriSpeech-100 and on synthetic long-form clips (with and without `--vad`) is in `docs/ternary.md`. + +| Model | Variant | Converter flags | Size | Notes | +|---|---|---|---:|---| +| parakeet-ultra | F16 | `--dtype f16` | 1441.9 MB | ordinary v3-shaped GGUF | +| parakeet-redux | packed ternary | `--ternary keep` | 213.3 MB | CPU only, no streaming, see `docs/ternary.md` | +| parakeet-redux | dequantized F16 | `--ternary dequant --dtype f16` | 1441.9 MB | runs on any backend | + +Expected file set for a release (names as used by the tests and docs): +`parakeet-ultra-f16.gguf`, `parakeet-redux-ternary.gguf` (`--ternary keep`, packed) and +`parakeet-redux-f16.gguf` (`--ternary dequant --dtype f16`). Use it for any GPU backend and for CPUs that would run the packed +file on the slow scalar kernel: MSVC builds, Windows on ARM and aarch64 without dotprod. x86-64 with AVX2 +or AVX-512 VNNI and aarch64 with dotprod get SIMD kernels for the packed file. The packed file also keeps the +original packed tensors resident next to the repacked planes, so its memory use is more than 213 MB. +Ultra Q8_0 (`--dtype q8_0`, 941.5 MB) is measured in `docs/ternary.md` but is optional. + +Both carry the `parakeet.vad.*` KVs and `vad_head.*` tensors unless converted +with `--vad drop`. + --- ## Notes diff --git a/scripts/compare_hyps.py b/scripts/compare_hyps.py new file mode 100644 index 0000000..44f112e --- /dev/null +++ b/scripts/compare_hyps.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +"""Compare two directories of bench-style hypothesis JSONs (/.json). + +Prints, per language, how many utterances are identical after normalize() and the +WER of A scored against B as the reference (A = parakeet.cpp, B = HF reference). + + python3 scripts/compare_hyps.py ours_dir hf_dir +""" +import json +import pathlib +import re +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from asr_metrics import _edit_distance, normalize # noqa: E402 + +def clean(t): + # The HF side prints the piece where parakeet.cpp drops it; ignore that + # one decoding-detail difference so it does not count as an ASR difference. + return re.sub(r"", " ", t) + + +a, b = map(pathlib.Path, sys.argv[1:3]) +tot_e = tot_w = tot_same = tot_n = 0 +for ja in sorted(a.glob("*.json")): + jb = b / ja.name + if not jb.exists(): + continue + fa = {f["path"]: f["text"] for f in json.loads(ja.read_text())["files"]} + fb = {f["path"]: f["text"] for f in json.loads(jb.read_text())["files"]} + e = w = same = 0 + for p in fb: + x, y = normalize(clean(fa[p])).split(), normalize(clean(fb[p])).split() + e += _edit_distance(y, x) + w += len(y) + same += x == y + if x != y: + print(f" DIFF {ja.stem} {pathlib.Path(p).name}\n ours: {fa[p]}\n hf: {fb[p]}") + print(f"{ja.stem}: identical {same}/{len(fb)} WER(ours vs hf) {100.0 * e / w:.2f}% ({e}/{w} words)") + tot_e += e; tot_w += w; tot_same += same; tot_n += len(fb) +print(f"TOTAL: identical {tot_same}/{tot_n} WER(ours vs hf) {100.0 * tot_e / tot_w:.2f}% ({tot_e}/{tot_w} words)") diff --git a/scripts/convert_hf_parakeet_to_gguf.py b/scripts/convert_hf_parakeet_to_gguf.py new file mode 100644 index 0000000..552a396 --- /dev/null +++ b/scripts/convert_hf_parakeet_to_gguf.py @@ -0,0 +1,312 @@ +#!/usr/bin/env python3 +"""Convert an HF-format Parakeet TDT checkpoint (safetensors) to GGUF. + +Targets checkpoints that follow the transformers ``ParakeetForTDT`` layout and +are derived from ``nvidia/parakeet-tdt-0.6b-v3``, for example +``moondream/parakeet-ultra`` (post-trained, F16) and ``moondream/parakeet-redux`` +(ternary encoder, packed ``qweight`` + ``scales``). + +There is no NeMo dependency. The HF repos carry neither the mel filterbank and +window nor a SentencePiece vocab, so the KV metadata, the tokenizer pieces and +the two featurizer buffers are taken from a GGUF of the teacher model that was +produced by ``convert_parakeet_to_gguf.py`` (``--template``). The HF +``config.json`` is cross-checked against the template and the run aborts on any +mismatch, so a checkpoint that is not v3-shaped cannot be converted silently. + +HF tensor names are mapped back to the verbatim NeMo names the C++ loader +expects (the inverse of transformers' ``convert_nemo_to_hf.py``). Ternary +weights are dequantized by default (``w = scales[row, col // group] * (code - 1)``), +so the engine sees ordinary linear weights and needs no new kernel. + +The optional ``vad_head.*`` tensors are written as F32 together with the +``parakeet.vad.*`` KVs unless ``--vad drop`` is given. With ``--ternary keep`` +the ternary linears stay packed (see docs/ternary.md) instead of dequantized. +""" +import argparse +import json +import pathlib +import re +import struct +import sys + +import numpy as np + +try: + import gguf +except ImportError as e: # pragma: no cover - env guard + print(f"converter: missing dependency 'gguf': {e}", file=sys.stderr) + sys.exit(2) + +# HF name -> NeMo name. Order matters only where patterns could overlap; every +# rule is anchored, so the first match wins. +_RENAMES = [ + (r"^encoder\.subsampling\.layers\.(\d+)\.(weight|bias)$", r"encoder.pre_encode.conv.\1.\2"), + (r"^encoder\.subsampling\.linear\.(weight|bias)$", r"encoder.pre_encode.out.\1"), + (r"^encoder\.layers\.(\d+)\.conv\.norm\.", r"encoder.layers.\1.conv.batch_norm."), + (r"^encoder\.layers\.(\d+)\.self_attn\.q_proj\.", r"encoder.layers.\1.self_attn.linear_q."), + (r"^encoder\.layers\.(\d+)\.self_attn\.k_proj\.", r"encoder.layers.\1.self_attn.linear_k."), + (r"^encoder\.layers\.(\d+)\.self_attn\.v_proj\.", r"encoder.layers.\1.self_attn.linear_v."), + (r"^encoder\.layers\.(\d+)\.self_attn\.o_proj\.", r"encoder.layers.\1.self_attn.linear_out."), + (r"^encoder\.layers\.(\d+)\.self_attn\.relative_k_proj\.", r"encoder.layers.\1.self_attn.linear_pos."), + (r"^encoder\.layers\.(\d+)\.self_attn\.bias_([uv])$", r"encoder.layers.\1.self_attn.pos_bias_\2"), + (r"^decoder\.embedding\.", "decoder.prediction.embed."), + (r"^decoder\.lstm\.", "decoder.prediction.dec_rnn.lstm."), + (r"^encoder_projector\.", "joint.enc."), + (r"^decoder\.decoder_projector\.", "joint.pred."), + (r"^joint\.head\.", "joint.joint_net.2."), +] +_RENAMES = [(re.compile(a), b) for a, b in _RENAMES] + +# Unchanged names (norms, FFN, conv, pointwise convs, ...) pass through. +_PASSTHROUGH = re.compile(r"^encoder\.layers\.\d+\.") + +# Linear weights that go through ggml_mul_mat and may be quantized. Mirrors the +# ASR entries of _QUANTIZABLE_PATTERNS in convert_parakeet_to_gguf.py and the +# policy in docs/quantization.md. +_QUANTIZABLE = [re.compile(p) for p in ( + r"^encoder\.layers\.\d+\.feed_forward[12]\.linear[12]\.weight$", + r"^encoder\.layers\.\d+\.self_attn\.linear_(q|k|v|out|pos)\.weight$", + r"^encoder\.pre_encode\.out\.weight$", + r"^joint\.enc\.weight$", + r"^joint\.pred\.weight$", +)] + +_VAD = re.compile(r"^vad_head\.") + + +def hf_to_nemo(name): + for rx, rep in _RENAMES: + if rx.search(name): + return rx.sub(rep, name) + if _PASSTHROUGH.match(name): + return name + raise KeyError(f"no NeMo name for HF tensor {name!r}") + + +def read_safetensors(path): + """Yield (name, numpy array or (dtype, shape, bytes)) without needing torch.""" + dt = {"F32": np.float32, "F16": np.float16, "I64": np.int64, "U8": np.uint8} + with open(path, "rb") as f: + (n,) = struct.unpack(" 2: + raise ValueError("ternary payload has a byte >= 3**5") + s = np.repeat(scales.astype(np.float32), group, axis=1)[:, :in_features] + return s * (codes.astype(np.float32) - 1.0) + + +def load_template(path): + r = gguf.GGUFReader(path) + kv = {} + for k, f in r.fields.items(): + if k.startswith("GGUF."): + continue + kv[k] = f + tensors = {t.name: t for t in r.tensors} + return r, kv, tensors + + +def field_value(f): + return f.contents() + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--hf", required=True, help="local dir holding model.safetensors + config.json (+ ternary.json)") + ap.add_argument("--template", required=True, + help="GGUF of nvidia/parakeet-tdt-0.6b-v3 from convert_parakeet_to_gguf.py (KV, vocab, mel buffers)") + ap.add_argument("--output", required=True) + ap.add_argument("--dtype", choices=["f32", "f16", "q8_0"], default="f32") + ap.add_argument("--name", default=None, help="general.name (default: the --hf directory name)") + ap.add_argument("--ternary", choices=["dequant", "keep"], default="dequant", + help="keep: store ternary linears packed as .qweight + .scales (needs the native ternary kernel); dequant: ordinary weights") + ap.add_argument("--vad", choices=["keep", "drop"], default="keep", + help="keep the vad_head.* tensors and parakeet.vad.* KVs when the checkpoint has them") + args = ap.parse_args() + + src = pathlib.Path(args.hf) + cfg = json.load(open(src / "config.json")) + tern_path = src / "ternary.json" + tern = json.load(open(tern_path)) if tern_path.exists() else None + + _, tkv, ttensors = load_template(args.template) + tv = {k: field_value(f) for k, f in tkv.items() if "pieces" not in k} + + # --- cross-check HF config against the template ------------------------- + enc = cfg["encoder_config"] + checks = { + "parakeet.encoder.d_model": enc["hidden_size"], + "parakeet.encoder.n_layers": enc["num_hidden_layers"], + "parakeet.encoder.n_heads": enc["num_attention_heads"], + "parakeet.encoder.ff_dim": enc["intermediate_size"], + "parakeet.encoder.conv_kernel": enc["conv_kernel_size"], + "parakeet.encoder.feat_in": enc["num_mel_bins"], + "parakeet.encoder.subsampling_factor": enc["subsampling_factor"], + "parakeet.encoder.subsampling_conv_channels": enc["subsampling_conv_channels"], + "parakeet.decoder.pred_hidden": cfg["decoder_hidden_size"], + "parakeet.decoder.pred_rnn_layers": cfg["num_decoder_layers"], + "parakeet.vocab_size": cfg["vocab_size"] - 1, + "parakeet.blank_id": cfg["blank_token_id"], + } + for k, want in checks.items(): + if int(tv[k]) != int(want): + sys.exit(f"config mismatch vs template: {k} template={tv[k]} hf={want}") + if [int(x) for x in tv["parakeet.tdt.durations"]] != cfg["durations"]: + sys.exit("config mismatch vs template: tdt durations") + + # --- KV: copy the template, HF checkpoints are TDT-only (no CTC head) ---- + w = gguf.GGUFWriter(args.output, "parakeet") + w.add_string("general.name", args.name or src.name) + for k, f in tkv.items(): + if k in ("general.architecture", "general.name", "general.quantization_version", + "general.file_type"): + continue + if k == "parakeet.arch": + w.add_string(k, "tdt") + elif k == "parakeet.tokenizer.pieces": + w.add_array(k, [str(p) for p in f.contents()]) + else: + _copy_kv(w, k, f) + + # --- tensors ------------------------------------------------------------ + keep_tern = args.ternary == "keep" + if keep_tern and not tern: + sys.exit("--ternary keep needs ternary.json next to model.safetensors") + tens = {} # float32 tensors; allowlisted linears are quantized per --dtype + raw = {} # written as-is: ternary qweight (int8 view) and scales (f16), vad_head.* + qmods = {m["name"]: m for m in (tern["quantized_modules"] if tern else [])} + group = (tern or {}).get("quant", {}).get("group_size", 0) + pending_q = {} + for name, arr in read_safetensors(src / "model.safetensors"): + if arr.ndim == 0: + continue + if _VAD.match(name): + if args.vad == "keep": + raw[name] = np.ascontiguousarray(arr, dtype=np.float32) + continue + if name.endswith(".qweight"): + pending_q.setdefault(name[:-8], {})["q"] = arr + continue + if name.endswith(".scales"): + pending_q.setdefault(name[:-7], {})["s"] = arr + continue + tens[hf_to_nemo(name)] = np.ascontiguousarray(arr, dtype=np.float32) + + for mod, d in pending_q.items(): + m = qmods[mod] + g = m["group_size"] or group + base = hf_to_nemo(mod + ".weight")[: -len(".weight")] + if keep_tern: + if g != 128 or m["in_features"] % g: + sys.exit(f"--ternary keep needs group 128 and in_features % 128 == 0 " + f"(got group {g}, in {m['in_features']} for {mod})") + raw[base + ".qweight"] = np.ascontiguousarray(d["q"]).view(np.int8) + raw[base + ".scales"] = np.ascontiguousarray(d["s"], dtype=np.float16) + continue + wt = unpack_ternary(d["q"], d["s"], m["in_features"], g) + if m.get("as_conv1d") or mod.rsplit(".", 1)[-1] in ("pointwise_conv1", "pointwise_conv2"): + wt = wt[:, :, None] # NeMo stores 1x1 convs as [out, in, 1] + tens[base + ".weight"] = np.ascontiguousarray(wt, dtype=np.float32) + + # mel featurizer buffers are not in the HF repo; lift them from the template + for k in ("preprocessor.featurizer.fb", "preprocessor.featurizer.window"): + tens[k] = np.ascontiguousarray(ttensors[k].data, dtype=np.float32) + + # The tensor set must equal the template's (minus the weights replaced by + # qweight/scales), with identical shapes. + replaced = {n[: -len(".qweight")] + ".weight" for n in raw if n.endswith(".qweight")} + expect = {n for n in ttensors + if not n.startswith("ctc_decoder.") and not n.startswith("decoder.decoder_layers")} - replaced + got = set(tens) + if expect != got: + sys.exit(f"tensor set differs from template: missing={sorted(expect - got)[:5]} " + f"extra={sorted(got - expect)[:5]}") + for n, a in tens.items(): + want = tuple(int(x) for x in ttensors[n].shape[::-1]) + if a.ndim > 1 and a.shape != want and n != "preprocessor.featurizer.fb": + sys.exit(f"shape mismatch {n}: hf={a.shape} template(numpy order)={want}") + + if keep_tern: + w.add_bool("parakeet.ternary.present", True) + w.add_uint32("parakeet.ternary.group_size", 128) + if any(n.startswith("vad_head.") for n in raw): + proj, ctxw = raw["vad_head.proj.weight"], raw["vad_head.ctx.weight"] + hop, sub, sr = (int(tv["parakeet.preprocessor.hop_length"]), + int(tv["parakeet.encoder.subsampling_factor"]), + int(tv["parakeet.preprocessor.sample_rate"])) + w.add_bool("parakeet.vad.present", True) + w.add_uint32("parakeet.vad.d_in", int(proj.shape[1])) + w.add_uint32("parakeet.vad.hidden", int(proj.shape[0])) + w.add_uint32("parakeet.vad.kernel", int(ctxw.shape[2])) + w.add_float32("parakeet.vad.frame_sec", hop * sub / sr) + + for n, a in raw.items(): + w.add_tensor(n, a) + + written = quantized = 0 + for n, a in tens.items(): + ne = list(a.shape[::-1]) + qt = None + if args.dtype != "f32" and any(rx.match(n) for rx in _QUANTIZABLE) and len(ne) >= 2 and ne[0] >= 32 and ne[1] >= 32: + if args.dtype == "f16": + qt = gguf.GGMLQuantizationType.F16 + elif ne[0] % 32 == 0: + qt = gguf.GGMLQuantizationType.Q8_0 + if qt is None: + w.add_tensor(n, a) + else: + qa = gguf.quantize(a, qt) + w.add_tensor(n, qa, raw_shape=qa.shape, raw_dtype=qt) + quantized += 1 + written += 1 + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print(f"wrote {args.output}: arch=tdt tensors={written + len(raw)} dtype={args.dtype} " + f"quantized={quantized} ternary={'keep' if keep_tern else 'dequant'} raw={len(raw)}") + + +def _copy_kv(w, key, f): + """Re-emit one template KV with its original GGUF type.""" + t = f.types + v = f.contents() + T = gguf.GGUFValueType + if t[0] == T.ARRAY: + w.add_array(key, list(v)) + return + {T.UINT32: w.add_uint32, T.INT32: w.add_int32, T.FLOAT32: w.add_float32, + T.BOOL: w.add_bool, T.STRING: w.add_string, T.UINT64: w.add_uint64, + T.INT64: w.add_int64, T.UINT8: w.add_uint8, T.UINT16: w.add_uint16}[t[0]](key, v) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_longform_talks.py b/scripts/eval_longform_talks.py new file mode 100644 index 0000000..bb26a99 --- /dev/null +++ b/scripts/eval_longform_talks.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Plain single-pass vs --vad WER on full-length talks (.wav + .txt). + +Like eval_vad_longform.py, but keeps every hypothesis, records wall time and peak +RSS (via /usr/bin/time -v) per run, and survives a failed single pass (a crash or +an out-of-memory kill is recorded as FAIL instead of aborting the sweep). +Existing hypothesis files are reused. Scoring uses asr_metrics.normalize on both +sides (not the Open ASR Leaderboard normalizer). + + python3 scripts/eval_longform_talks.py --model ultra.gguf --dir /tmp/val/ted --save /tmp/val/ted_res/ultra +""" +import argparse +import json +import pathlib +import re +import subprocess +import sys +import time + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from asr_metrics import wer # noqa: E402 + + +def run(cli, model, wav, extra, threads, timeout): + t0 = time.time() + try: + r = subprocess.run(["/usr/bin/time", "-v", cli, "transcribe", "--model", model, "--input", str(wav), + "--decoder", "tdt", "--threads", str(threads), *extra], + capture_output=True, text=True, timeout=timeout) + except subprocess.TimeoutExpired: + return None, time.time() - t0, 0, "timeout" + m = re.search(r"Maximum resident set size \(kbytes\): (\d+)", r.stderr) + rss = int(m.group(1)) / 1e6 if m else 0.0 + if r.returncode != 0: + return None, time.time() - t0, rss, f"exit {r.returncode}: {r.stderr.strip().splitlines()[-1] if r.stderr.strip() else ''}" + return r.stdout.strip(), time.time() - t0, rss, "" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--cli", default="build/examples/cli/parakeet-cli") + ap.add_argument("--model", required=True) + ap.add_argument("--dir", required=True) + ap.add_argument("--save", required=True) + ap.add_argument("--threads", type=int, default=8) + ap.add_argument("--timeout", type=int, default=3600) + a = ap.parse_args() + save = pathlib.Path(a.save) + save.mkdir(parents=True, exist_ok=True) + rows = [] + for wav in sorted(pathlib.Path(a.dir).glob("*.wav")): + ref = wav.with_suffix(".txt").read_text(encoding="utf-8") + row = {"talk": wav.stem} + for mode, extra in (("plain", []), ("vad", ["--vad"])): + hp, mp = save / f"{wav.stem}.{mode}.txt", save / f"{wav.stem}.{mode}.json" + if mp.exists(): + row[mode] = json.loads(mp.read_text()) + continue + hyp, secs, rss, err = run(a.cli, a.model, wav, extra, a.threads, a.timeout) + res = {"wall_s": round(secs, 1), "rss_gb": round(rss, 2), "err": err, + "wer": None if hyp is None else round(100 * wer(ref, hyp), 3)} + if hyp is not None: + hp.write_text(hyp + "\n", encoding="utf-8") + mp.write_text(json.dumps(res)) + row[mode] = res + print(wav.stem, mode, res, flush=True) + rows.append(row) + ok = [r for r in rows if r["plain"]["wer"] is not None and r["vad"]["wer"] is not None] + print("talk\tplain_wer\tvad_wer\tplain_s\tvad_s\tplain_rss_gb\tvad_rss_gb") + for r in rows: + p, v = r["plain"], r["vad"] + print(f"{r['talk']}\t{p['wer'] if p['wer'] is not None else 'FAIL'}\t{v['wer'] if v['wer'] is not None else 'FAIL'}" + f"\t{p['wall_s']}\t{v['wall_s']}\t{p['rss_gb']}\t{v['rss_gb']}") + for mode in ("plain", "vad"): + v = [r[mode]["wer"] for r in rows if r[mode]["wer"] is not None] + if v: + print(f"MEAN {mode}: {sum(v) / len(v):.2f} over {len(v)}/{len(rows)} talks") + if ok: + print(f"MEAN on the {len(ok)} talks where both ran: plain {sum(r['plain']['wer'] for r in ok) / len(ok):.2f} " + f"vad {sum(r['vad']['wer'] for r in ok) / len(ok):.2f}") + +if __name__ == "__main__": + main() diff --git a/scripts/eval_manifest_wer.py b/scripts/eval_manifest_wer.py new file mode 100644 index 0000000..ea41cc7 --- /dev/null +++ b/scripts/eval_manifest_wer.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +"""Corpus WER of one or more GGUF models over per-language manifests. + +Each manifest is a TSV of `wav_pathreference`. For every (model, manifest) +pair this runs `parakeet-cli bench --json`, keeps the hypotheses in +//.json (reused if present, so reruns are +cheap) and scores with scripts/asr_metrics.normalize on BOTH sides. That is not +the Open ASR Leaderboard normalizer, so absolute numbers are not comparable with +published model cards; compare models against each other on the same subset. + +WER per manifest is a corpus WER (total edits / total reference words). The +summary line is the unweighted mean over manifests (macro average). + + python3 scripts/eval_manifest_wer.py --out /tmp/val/res \ + --model v3=models/v3.gguf --model ultra=models/ultra.gguf \ + /tmp/val/fleurs/en_us/manifest.tsv /tmp/val/fleurs/de_de/manifest.tsv +""" +import argparse +import json +import pathlib +import subprocess +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from asr_metrics import _edit_distance, normalize # noqa: E402 + + +def score(manifest, doc): + refs = {} + for line in pathlib.Path(manifest).read_text(encoding="utf-8").splitlines(): + if line.strip(): + p, r = line.split("\t", 1) + refs[p] = r + edits = words = 0 + for f in doc["files"]: + r = normalize(refs[f["path"]]).split() + h = normalize(f["text"]).split() + edits += _edit_distance(r, h) + words += len(r) + audio = sum(f["audio_sec"] for f in doc["files"]) + proc = sum(f["proc_ms"] for f in doc["files"]) / 1000.0 + return edits, words, audio, proc, len(doc["files"]) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("manifests", nargs="+") + ap.add_argument("--model", action="append", required=True, help="name=path.gguf (repeatable)") + ap.add_argument("--out", required=True) + ap.add_argument("--cli", default="build/examples/cli/parakeet-cli") + ap.add_argument("--threads", type=int, default=8) + ap.add_argument("--decoder", default="tdt") + a = ap.parse_args() + out = pathlib.Path(a.out) + table = {} + for spec in a.model: + name, path = spec.split("=", 1) + (out / name).mkdir(parents=True, exist_ok=True) + for m in a.manifests: + lang = pathlib.Path(m).parent.name + js = out / name / f"{lang}.json" + if not js.exists(): + subprocess.run([a.cli, "bench", "--model", path, "--manifest", m, "--decoder", + a.decoder, "--threads", str(a.threads), "--json", str(js)], + check=True, stdout=subprocess.DEVNULL) + e, w, au, pr, n = score(m, json.loads(js.read_text())) + table[(name, lang)] = (100.0 * e / w, w, au, n) + names = [s.split("=", 1)[0] for s in a.model] + langs = sorted({l for _, l in table}) + print("lang\t" + "\t".join(names)) + for l in langs: + print(l + "\t" + "\t".join(f"{table[(n, l)][0]:.2f}" for n in names)) + print("MEAN\t" + "\t".join(f"{sum(table[(n, l)][0] for l in langs) / len(langs):.2f}" for n in names)) + n0 = names[0] + print(f"utterances={sum(table[(n0, l)][3] for l in langs)} " + f"ref_words={sum(table[(n0, l)][1] for l in langs)} " + f"audio_sec={sum(table[(n0, l)][2] for l in langs):.1f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_vad_longform.py b/scripts/eval_vad_longform.py new file mode 100644 index 0000000..0a8697f --- /dev/null +++ b/scripts/eval_vad_longform.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""WER of plain single-pass vs --vad transcription on long clips with references. + +Use --skip-plain to reuse a plain-pass result when sweeping VAD parameters. +""" +import argparse +import pathlib +import subprocess +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from asr_metrics import wer # noqa: E402 + + +def run(cmd_prefix, cli, model, wav, extra, threads): + try: + r = subprocess.run([*cmd_prefix, cli, "transcribe", "--model", model, "--input", wav, + "--decoder", "tdt", "--threads", str(threads), *extra], + capture_output=True, text=True, check=True) + except subprocess.CalledProcessError as e: + print(f"parakeet-cli failed (exit {e.returncode}) on {wav}:\n{e.stderr}", file=sys.stderr) + sys.exit(1) + return r.stdout.strip() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--cli", default="build/examples/cli/parakeet-cli") + ap.add_argument("--model", required=True) + ap.add_argument("--dir", required=True) + ap.add_argument("--glob", default="*.wav") + ap.add_argument("--threads", type=int, default=8) + ap.add_argument("--prefix", default="", help="command prefix, e.g. 'taskset -c 0-7'") + ap.add_argument("--skip-plain", action="store_true") + ap.add_argument("--save", help="directory to keep the raw transcripts") + ap.add_argument("--vad-arg", action="append", default=[], + help="extra args for the --vad run, e.g. --vad-arg=--vad-threshold=0.4 (repeatable)") + args = ap.parse_args() + extra = [a for v in args.vad_arg for a in v.split("=", 1)] + prefix = args.prefix.split() + save = pathlib.Path(args.save) if args.save else None + if save: + save.mkdir(parents=True, exist_ok=True) + tot_plain = tot_vad = n = 0.0 + wavs = sorted(pathlib.Path(args.dir).glob(args.glob)) + if not wavs: + print(f"no files match {args.glob} in {args.dir}", file=sys.stderr) + sys.exit(2) + for wav in wavs: + ref = wav.with_suffix(".txt").read_text() + wp = float("nan") + if not args.skip_plain: + plain = run(prefix, args.cli, args.model, str(wav), [], args.threads) + wp = wer(ref, plain) + if save: + (save / (wav.stem + ".plain.txt")).write_text(plain + "\n") + vad = run(prefix, args.cli, args.model, str(wav), ["--vad", *extra], args.threads) + wv = wer(ref, vad) + if save: + (save / (wav.stem + ".vad.txt")).write_text(vad + "\n") + print(f"{wav.name:32s} plain {wp*100:5.2f}% vad {wv*100:5.2f}%") + tot_plain += wp; tot_vad += wv; n += 1 + print(f"{'mean':32s} plain {tot_plain/n*100:5.2f}% vad {tot_vad/n*100:5.2f}%") + + +if __name__ == "__main__": + main() diff --git a/scripts/fetch_fleurs_subset.py b/scripts/fetch_fleurs_subset.py new file mode 100644 index 0000000..6714e2a --- /dev/null +++ b/scripts/fetch_fleurs_subset.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +"""Stream the first N test utterances of FLEURS languages into 16 kHz mono wavs. + +For each FLEURS config id it writes //NNNN.wav and //manifest.tsv +(`wav_pathreference`, reference = the dataset's `raw_transcription`, the +natural-cased text with punctuation; scoring normalizes it). Utterances are the +first N of the test split in dataset (parquet) order, no shuffling or filtering. + +Only public, ungated data is used and no token is needed. Audio is streamed with +decode=False and decoded here with soundfile, so torchcodec is not required. + + python3 scripts/fetch_fleurs_subset.py --out /tmp/val/fleurs --langs en_us,de_de --n 50 +""" +import argparse +import io +import pathlib + +import librosa +import soundfile as sf +from datasets import Audio, load_dataset + +ALL = ("bg_bg hr_hr cs_cz da_dk nl_nl en_us et_ee fi_fi fr_fr de_de el_gr hu_hu it_it " + "lv_lv lt_lt mt_mt pl_pl pt_br ro_ro ru_ru sk_sk sl_si es_419 sv_se uk_ua").split() + + +def fetch(cfg, out, n): + d = pathlib.Path(out) / cfg + d.mkdir(parents=True, exist_ok=True) + ds = load_dataset("google/fleurs", cfg, split="test", streaming=True) + ds = ds.cast_column("audio", Audio(decode=False)) + rows, secs = [], 0.0 + for i, ex in enumerate(ds): + if i >= n: + break + y, sr = sf.read(io.BytesIO(ex["audio"]["bytes"]), dtype="float32", always_2d=False) + if y.ndim > 1: + y = y.mean(axis=1) + if sr != 16000: + y = librosa.resample(y, orig_sr=sr, target_sr=16000) + p = d / f"{i:04d}.wav" + sf.write(str(p), y, 16000, subtype="PCM_16") + secs += len(y) / 16000 + ref = " ".join(ex["raw_transcription"].split()) + rows.append(f"{p}\t{ref}") + (d / "manifest.tsv").write_text("\n".join(rows) + "\n", encoding="utf-8") + return len(rows), secs + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--out", required=True) + ap.add_argument("--langs", default=",".join(ALL)) + ap.add_argument("--n", type=int, default=50) + a = ap.parse_args() + for cfg in a.langs.split(","): + k, s = fetch(cfg, a.out, a.n) + print(f"{cfg}\t{k} utts\t{s:.1f} s", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/fetch_tedlium_longform.py b/scripts/fetch_tedlium_longform.py new file mode 100644 index 0000000..22693df --- /dev/null +++ b/scripts/fetch_tedlium_longform.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Stream distil-whisper/tedlium-long-form (test split, 11 full talks) to disk. + +Writes /.wav (16 kHz mono int16) and /.txt (the reference +with tags such as removed and whitespace collapsed). Public, ungated, no +token needed. Audio is streamed with decode=False and decoded with soundfile. + + python3 scripts/fetch_tedlium_longform.py --out /tmp/val/tedlium +""" +import argparse +import io +import pathlib +import re + +import librosa +import soundfile as sf +from datasets import Audio, load_dataset + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--out", required=True) + ap.add_argument("--only", help="write only the talk whose name contains this text (e.g. BillGates)") + a = ap.parse_args() + out = pathlib.Path(a.out) + out.mkdir(parents=True, exist_ok=True) + ds = load_dataset("distil-whisper/tedlium-long-form", split="test", streaming=True) + ds = ds.cast_column("audio", Audio(decode=False)) + for i, ex in enumerate(ds): + name = re.sub(r"[^A-Za-z0-9_-]", "", pathlib.Path(ex["audio"]["path"] or f"talk{i}").stem) or f"talk{i}" + if a.only and a.only not in name: + continue + y, sr = sf.read(io.BytesIO(ex["audio"]["bytes"]), dtype="float32") + if y.ndim > 1: + y = y.mean(axis=1) + if sr != 16000: + y = librosa.resample(y, orig_sr=sr, target_sr=16000) + sf.write(str(out / f"{name}.wav"), y, 16000, subtype="PCM_16") + text = " ".join(re.sub(r"<[^>]*>", " ", ex["text"]).split()) + (out / f"{name}.txt").write_text(text + "\n", encoding="utf-8") + print(f"{name}\t{len(y) / 16000:.1f} s\t{len(text.split())} words", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/hf_reference_transcribe.py b/scripts/hf_reference_transcribe.py new file mode 100644 index 0000000..6096a54 --- /dev/null +++ b/scripts/hf_reference_transcribe.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Transcribe manifests with the transformers ParakeetForTDT reference (CPU). + +Used to check parakeet.cpp against an independent implementation of a Parakeet +derivative that only ships in HF format (e.g. moondream/parakeet-ultra). The +HF checkpoint has no preprocessor or tokenizer files, so this builds a +ParakeetFeatureExtractor with the checkpoint's mel count and decodes token ids +with the sentencepiece piece table stored in a parakeet.cpp GGUF of the same +model (`parakeet.tokenizer.pieces`; a leading U+2581 is a word start). + +ParakeetForTDT is not in every transformers release; point PYTHONPATH at a +transformers source tree that has it if the installed one does not. + + PYTHONPATH=/src python3 scripts/hf_reference_transcribe.py \ + --hf-dir --gguf ultra-f16.gguf --out hyp/ en_us/manifest.tsv ... + +Writes /.json in the same {"files":[{path,text}]} shape +as `parakeet-cli bench --json`, so scripts/eval_manifest_wer.py can score it. +""" +import argparse +import json +import pathlib + +import gguf +import numpy as np +import soundfile as sf +import torch +from transformers import ParakeetFeatureExtractor, ParakeetForTDT + + +def pieces_from_gguf(path): + f = gguf.GGUFReader(path).fields["parakeet.tokenizer.pieces"] + return [bytes(f.parts[i]).decode("utf-8") for i in f.data] + + +def detok(ids, pieces, blank): + return "".join(pieces[i] for i in ids if i != blank and i < len(pieces)).replace("▁", " ").strip() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("manifests", nargs="+") + ap.add_argument("--hf-dir", required=True) + ap.add_argument("--gguf", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--threads", type=int, default=8) + a = ap.parse_args() + torch.set_num_threads(a.threads) + model = ParakeetForTDT.from_pretrained(a.hf_dir, torch_dtype=torch.float32).eval() + fe = ParakeetFeatureExtractor(feature_size=model.config.encoder_config.num_mel_bins, sampling_rate=16000) + pieces = pieces_from_gguf(a.gguf) + blank = model.config.blank_token_id + out = pathlib.Path(a.out) + out.mkdir(parents=True, exist_ok=True) + for m in a.manifests: + files = [] + for line in pathlib.Path(m).read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + p = line.split("\t", 1)[0] + y, sr = sf.read(p, dtype="float32") + assert sr == 16000 + feats = fe(y, sampling_rate=16000, return_tensors="pt", return_attention_mask=True) + with torch.no_grad(): + seq = model.generate(**feats, decoder_start_token_id=blank, + suppress_tokens=list(range(model.config.vocab_size, model.config.vocab_size + len(model.config.durations)))) + seq = seq.sequences if hasattr(seq, "sequences") else seq + files.append({"path": p, "text": detok(seq[0].tolist(), pieces, blank), "audio_sec": len(y) / 16000, "proc_ms": 0}) + (out / f"{pathlib.Path(m).parent.name}.json").write_text(json.dumps({"files": files}, ensure_ascii=False)) + print(m, len(files), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/make_longform.py b/scripts/make_longform.py new file mode 100644 index 0000000..a00e7a4 --- /dev/null +++ b/scripts/make_longform.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Join LibriSpeech benchmark utterances into long clips with a known reference. + +Each output clip is N utterances separated by a gap of low-level Gaussian noise +(about -55 dBFS, fixed seed) rather than digital zeros: digital silence distorts +the per-feature mel normalization and would not resemble real recordings. +Clips are written as /longform__.wav with a matching .txt holding +the joined reference. A gap of 0 joins the utterances with no inserted audio. +""" +import argparse +import pathlib +import wave + +import numpy as np + +NOISE_STD = 58.0 # int16 units, about -55 dBFS + + +def read(path): + with wave.open(path) as w: + assert w.getframerate() == 16000 and w.getnchannels() == 1 and w.getsampwidth() == 2, path + return np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--manifest", default="benchmarks/librispeech_manifest.tsv") + ap.add_argument("--out", required=True) + ap.add_argument("--per-clip", type=int, default=30) + ap.add_argument("--gap", type=float, default=0.45, help="noise gap between utterances, seconds (0 = none)") + ap.add_argument("--seed", type=int, default=1234) + args = ap.parse_args() + rows = [] + for line in open(args.manifest): + if not line.strip() or line.startswith("#"): + continue + path, text = line.rstrip("\n").split("\t", 1) + rows.append((path, text)) + out = pathlib.Path(args.out) + out.mkdir(parents=True, exist_ok=True) + rng = np.random.default_rng(args.seed) + n_gap = int(args.gap * 16000) + tag = f"{args.gap:.2f}".replace(".", "p") + for i in range(0, len(rows) - args.per_clip + 1, args.per_clip): + chunk = rows[i:i + args.per_clip] + parts = [] + for p, _ in chunk: + parts.append(read(p)) + if n_gap: + parts.append(np.clip(rng.normal(0, NOISE_STD, n_gap), -32768, 32767).astype(np.int16)) + audio = np.concatenate(parts) + stem = out / f"longform_{tag}_{i // args.per_clip}" + with wave.open(str(stem) + ".wav", "wb") as w: + w.setnchannels(1) + w.setsampwidth(2) + w.setframerate(16000) + w.writeframes(audio.tobytes()) + pathlib.Path(str(stem) + ".txt").write_text(" ".join(t for _, t in chunk) + "\n") + print(f"{stem}.wav {len(audio) / 16000:.1f} s {len(chunk)} utterances") + + +if __name__ == "__main__": + main() diff --git a/src/conformer.cpp b/src/conformer.cpp index 0622f05..293661f 100644 --- a/src/conformer.cpp +++ b/src/conformer.cpp @@ -2,10 +2,12 @@ #include "relpos_attention.hpp" #include "ggml_graph.hpp" #include "backend.hpp" +#include "ternary.hpp" #include "ggml.h" #include #include #include +#include #include #include @@ -46,10 +48,15 @@ static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml, const int pad = (K - 1) / 2; // symmetric padding (offline model) // -- pointwise_conv1 (Conv1d d->2d, k=1): 1x1 conv == linear over channels. - ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight"); - pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d] ggml_tensor* pw1b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv1.bias"); - ggml_tensor* y = ggml_mul_mat(ctx, pw1w, c); // [2d, T, B] + ggml_tensor* y; + if (has_ternary(ml, pre + "conv.pointwise_conv1")) { + y = ternary_linear(ctx, ml, pre + "conv.pointwise_conv1", c); // [2d, T, B] + } else { + ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight"); + pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d] + y = ggml_mul_mat(ctx, pw1w, c); // [2d, T, B] + } if (pw1b) y = ggml_add(ctx, y, pw1b); // -- GLU over channel dim (NeMo F.glu(x, dim=1)). y is [2D, T, B]; each half @@ -155,10 +162,15 @@ static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml, // -- SiLU (Swish), then pointwise_conv2 (Conv1d d->d, k=1). normed = ggml_silu(ctx, normed); - ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight"); - pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d] ggml_tensor* pw2b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv2.bias"); - ggml_tensor* cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T, B] + ggml_tensor* cout; + if (has_ternary(ml, pre + "conv.pointwise_conv2")) { + cout = ternary_linear(ctx, ml, pre + "conv.pointwise_conv2", normed); // [d, T, B] + } else { + ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight"); + pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d] + cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T, B] + } if (pw2b) cout = ggml_add(ctx, cout, pw2b); return cout; // [D, T, B]; this is layers[i].conv output } @@ -179,10 +191,15 @@ static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml, const int pad = (K - 1) / 2; // symmetric padding (offline model) // -- pointwise_conv1 (Conv1d d->2d, k=1): 1x1 conv == linear over channels. - ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight"); - pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d] ggml_tensor* pw1b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv1.bias"); - ggml_tensor* y = ggml_mul_mat(ctx, pw1w, c); // [2d, T] + ggml_tensor* y; + if (has_ternary(ml, pre + "conv.pointwise_conv1")) { + y = ternary_linear(ctx, ml, pre + "conv.pointwise_conv1", c); // [2d, T] + } else { + ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight"); + pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d] + y = ggml_mul_mat(ctx, pw1w, c); // [2d, T] + } if (pw1b) y = ggml_add(ctx, y, pw1b); // -- GLU over channel dim (NeMo F.glu(x, dim=1)). @@ -260,10 +277,15 @@ static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml, // -- SiLU (Swish), then pointwise_conv2 (Conv1d d->d, k=1). normed = ggml_silu(ctx, normed); - ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight"); - pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d] ggml_tensor* pw2b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv2.bias"); - ggml_tensor* cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T] + ggml_tensor* cout; + if (has_ternary(ml, pre + "conv.pointwise_conv2")) { + cout = ternary_linear(ctx, ml, pre + "conv.pointwise_conv2", normed); // [d, T] + } else { + ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight"); + pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d] + cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T] + } if (pw2b) cout = ggml_add(ctx, cout, pw2b); return cout; // [D, T] -> row-major [T, D]; this is layers[i].conv output } @@ -310,8 +332,9 @@ ggml_tensor* ConformerLayer::build_graph_batched(ggml_context* ctx, }; // nn.Linear: ggml weight ne = [in, out]. in ne [in, T, B] -> [out, T, B]. auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) { - ggml_tensor* W = clone_weight(ctx, ml, pre + nm + ".weight"); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); + ggml_tensor* y = has_ternary(ml, pre + nm) + ? ternary_linear(ctx, ml, pre + nm, in) + : ggml_mul_mat(ctx, clone_weight(ctx, ml, pre + nm + ".weight"), in); if (bias) { ggml_tensor* B = clone_weight_opt(ctx, ml, pre + nm + ".bias"); if (B) y = ggml_add(ctx, y, B); @@ -387,8 +410,9 @@ ggml_tensor* ConformerLayer::build_graph(ggml_context* ctx, ggml_tensor* xt, }; // nn.Linear: ggml weight ne = [in, out]. in ne [in, T] -> [out, T]. auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) { - ggml_tensor* W = clone_weight(ctx, ml, pre + nm + ".weight"); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); + ggml_tensor* y = has_ternary(ml, pre + nm) + ? ternary_linear(ctx, ml, pre + nm, in) + : ggml_mul_mat(ctx, clone_weight(ctx, ml, pre + nm + ".weight"), in); if (bias) { ggml_tensor* B = clone_weight_opt(ctx, ml, pre + nm + ".bias"); if (B) y = ggml_add(ctx, y, B); @@ -496,8 +520,14 @@ void ConformerLayer::forward_with_conv(const std::vector& x, int T, return y; }; auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) { - ggml_tensor* W = clone_weight(ctx, ml_, pre + nm + ".weight"); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); + ggml_tensor* y; + if (has_ternary(ml_, pre + nm)) { + y = ternary_linear(ctx, ml_, pre + nm, in); + } else { + if (!ml_.tensor(pre + nm + ".weight")) + throw std::runtime_error("missing encoder weight " + pre + nm + ".weight"); + y = ggml_mul_mat(ctx, clone_weight(ctx, ml_, pre + nm + ".weight"), in); + } if (bias) { ggml_tensor* B = clone_weight_opt(ctx, ml_, pre + nm + ".bias"); if (B) y = ggml_add(ctx, y, B); } return y; diff --git a/src/model.cpp b/src/model.cpp index 1912273..7644d76 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -1,10 +1,14 @@ #include "model.hpp" #include "audio_io.hpp" +#include "common.hpp" +#include "ternary.hpp" #include "mel.hpp" #include "mel_gpu.hpp" #include "encoder.hpp" #include "subsampling.hpp" +#include "vad_head.hpp" +#include "vad_segmenter.hpp" #include "ctc_decoder.hpp" #include "search.hpp" #include "tokenizer.hpp" @@ -20,7 +24,9 @@ #include "ggml_graph.hpp" #include +#include #include +#include #include #include @@ -50,6 +56,32 @@ std::unique_ptr Model::load(const std::string& gguf_path) { if (m->loader_.config().arch == "diarization") { return nullptr; } + // A GGUF whose packed tensors and ternary flag disagree would skip the GPU + // refusal and the validation below, so refuse it outright. + { + const std::string err = ternary_flag_consistency_error(m->loader_); + if (!err.empty()) { + PK_LOG("%s; refusing to load it", err.c_str()); + return nullptr; + } + } + // Packed ternary weights run on the CPU kernel only. Fail at load with a + // clear message instead of crashing inside a GPU graph. + if (m->loader_.config().ternary.present && + std::string(pk::global_backend().device_name()) != "cpu") { + PK_LOG("this GGUF holds packed ternary weights, which run on the CPU backend only; " + "re-convert with --ternary dequant to use a GPU backend, or set PARAKEET_DEVICE=cpu"); + return nullptr; + } + // Validate and repack every packed linear now so graph building never throws. + if (m->loader_.config().ternary.present) { + try { + ternary_prepare(m->loader_); + } catch (const std::exception& e) { + PK_LOG("invalid packed ternary GGUF: %s", e.what()); + return nullptr; + } + } // Give the weights a CPU backend buffer ONCE so graphs reference them // directly as leaves (zero per-call copy). Done at load (vs. lazily on first // clone_weight) so the cost is paid up front, not per utterance. @@ -57,6 +89,13 @@ std::unique_ptr Model::load(const std::string& gguf_path) { return m; } +const Model::DecoderObjects& Model::decoder_objects() const { + std::call_once(decoder_once_, [this]() { + decoder_ = std::make_unique(loader_); + }); + return *decoder_; +} + // Forward declarations: subsampling-tiling helpers are defined below (after the // batched staging helpers) but used by the single-clip transcribe entry points. static int safe_mel_window(const pk::ParakeetConfig& cfg); @@ -144,6 +183,7 @@ static EncodedAudio encode_16k(const ModelLoader& loader, // 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, + const PredictionNet* dpred, const Joint* djoint, const std::vector& enc_out, int d_model, int Tout, bool use_tdt) { const ParakeetConfig& cfg = loader.config(); @@ -152,8 +192,8 @@ static std::string decode_enc_out(const ModelLoader& loader, for (int t = 0; t < Tout; ++t) for (int c = 0; c < d_model; ++c) enc_row[(size_t)t * d_model + c] = enc_out[(size_t)c * Tout + t]; - PredictionNet pred(loader); - Joint joint(loader); + const PredictionNet& pred = *dpred; + const Joint& joint = *djoint; const int max_symbols = static_cast(cfg.max_symbols); std::vector ids; if (!cfg.tdt_durations.empty()) @@ -184,7 +224,10 @@ std::string Model::transcribe_16k(const std::vector& pcm16k, const bool use_tdt = (decoder == Decoder::kTDT) || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); - return decode_enc_out(loader_, encoded.channels_first, + return decode_enc_out(loader_, + use_tdt ? &decoder_objects().pred : nullptr, + use_tdt ? &decoder_objects().joint : nullptr, + encoded.channels_first, encoded.d_model, encoded.frames, use_tdt); } @@ -239,6 +282,117 @@ void Model::transcribe_16k_ctc_logits(const std::vector& pcm16k, T = Tout; } +std::vector Model::vad_probabilities(const std::vector& pcm16k, + const VadVariant* v) const { + const ParakeetConfig& cfg = loader_.config(); + if (!cfg.vad.present) throw std::runtime_error("model has no VAD head"); + 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); + } + Subsampling sub(loader_); + std::vector out; + int Tout = 0, d_model = 0, valid = 0; + const int tile = subsampling_tile_for(cfg, loader_, T); + if (tile > 0) sub.forward_tiled(feats, n_mels, T, tile, out, Tout, d_model, valid); + else sub.forward(feats, n_mels, T, out, Tout, d_model, valid); + if ((uint32_t)d_model != cfg.vad.d_in) + throw std::runtime_error("VAD head input width does not match the subsampler output"); + VadHead head(loader_); + return head.probabilities(out.data(), valid, v); +} + +namespace { + +// Slices the caller decodes: [first_sample, last_sample) of the 16 kHz PCM. +struct Slice { std::vector pcm; double start_sec; int start_frame; }; + +std::vector vad_slices(const Model& m, const std::vector& pcm16k, + const SegmenterOpts& opts_in) { + SegmenterOpts opts = opts_in; + opts.frame_sec = m.config().vad.frame_sec; + // Token frame offsets are in ENCODER frames (same formula as the JSON writer). + const ParakeetConfig& cfg = m.config(); + const double enc_frame_sec = + (double)cfg.hop_length * (double)cfg.subsampling_factor / (double)cfg.sample_rate; + if (!(enc_frame_sec > 0.0) || !std::isfinite(enc_frame_sec)) + throw std::runtime_error("invalid encoder frame size"); + const double total_sec = (double)pcm16k.size() / 16000.0; + const std::vector p = m.vad_probabilities(pcm16k); + const std::vector segs = segment_by_vad(p, total_sec, opts); + std::vector out; + const size_t n = pcm16k.size(); + auto at = [&](double sec) { + const long long v = std::llround(sec * 16000.0); + return (size_t)std::min(std::max(v, 0), (long long)n); + }; + for (size_t i = 0; i < segs.size(); ++i) { + // Consecutive slices share the exact boundary sample: a slice ends where + // the next one starts. + const size_t a = (i == 0) ? 0 : at(segs[i].start); + const size_t b = (i + 1 == segs.size()) ? n : std::max(a, at(segs[i + 1].start)); + Slice s; + s.pcm.assign(pcm16k.begin() + (std::ptrdiff_t)a, pcm16k.begin() + (std::ptrdiff_t)b); + if (s.pcm.size() < 3200) s.pcm.resize(3200, 0.0f); // 0.2 s minimum + s.start_sec = segs[i].start; + s.start_frame = (int)std::llround(segs[i].start / enc_frame_sec); + out.push_back(std::move(s)); + } + return out; +} + +} // namespace + +std::string Model::transcribe_pcm_vad(const std::vector& pcm, int sample_rate, + Decoder decoder, const std::string& target_lang, + const SegmenterOpts& opts) const { + if (!loader_.config().vad.present) throw std::runtime_error("model has no VAD head"); + const std::vector pcm16k = + sample_rate == 16000 ? pcm : resample_linear(pcm, sample_rate, 16000); + if ((double)pcm16k.size() / 16000.0 <= opts.max_seg_sec) + return transcribe_16k(pcm16k, decoder, target_lang); + const std::vector slices = vad_slices(*this, pcm16k, opts); + if (slices.empty()) return transcribe_16k(pcm16k, decoder, target_lang); + std::string text; + for (const Slice& s : slices) { + const std::string t = transcribe_16k(s.pcm, decoder, target_lang); + if (t.empty()) continue; + if (!text.empty()) text += ' '; + text += t; + } + return text; +} + +Transcription Model::transcribe_pcm_vad_with_timestamps(const std::vector& pcm, int sample_rate, + Decoder decoder, const std::string& target_lang, + const SegmenterOpts& opts) const { + if (!loader_.config().vad.present) throw std::runtime_error("model has no VAD head"); + const std::vector pcm16k = + sample_rate == 16000 ? pcm : resample_linear(pcm, sample_rate, 16000); + if ((double)pcm16k.size() / 16000.0 <= opts.max_seg_sec) + return transcribe_with_timestamps(pcm16k, 16000, decoder, target_lang); + const std::vector slices = vad_slices(*this, pcm16k, opts); + if (slices.empty()) return transcribe_with_timestamps(pcm16k, 16000, decoder, target_lang); + Transcription all; + for (const Slice& s : slices) { + Transcription t = transcribe_with_timestamps(s.pcm, 16000, decoder, target_lang); + for (Word& w : t.words) { w.start += (float)s.start_sec; w.end += (float)s.start_sec; } + for (TokenInfo& k : t.tokens) k.frame += s.start_frame; + if (!t.text.empty()) { + if (!all.text.empty()) all.text += ' '; + all.text += t.text; + } + all.words.insert(all.words.end(), t.words.begin(), t.words.end()); + all.tokens.insert(all.tokens.end(), t.tokens.begin(), t.tokens.end()); + } + return all; +} + // Max mel frames per encoder pass before the first subsampling conv output // (n_mels/2 * T/2 * conv_channels) approaches INT_MAX. ggml's CUDA unary (relu) // kernel indexes elements with int32, so a tensor > 2^31 elements crashes @@ -357,8 +511,8 @@ std::vector Model::transcribe_16k_batch( std::vector> encs; std::vector Ts; batch_enc_to_row_major(enc_outs, valid_Tout, d_model, encs, Ts); - PredictionNet pred(loader_); - Joint joint(loader_); + const PredictionNet& pred = decoder_objects().pred; + const Joint& joint = decoder_objects().joint; std::vector> ids; pk::transducer_greedy_batch(pred, joint, encs, Ts, d_model, cfg.tdt_durations, (int)cfg.blank_id, @@ -369,7 +523,10 @@ std::vector Model::transcribe_16k_batch( } else { // CTC stays per-item (no autoregressive decode to batch). for (int b = 0; b < mb.B; ++b) - outs[b] = decode_enc_out(loader_, enc_outs[b], d_model, valid_Tout[b], use_tdt); + outs[b] = decode_enc_out(loader_, + use_tdt ? &decoder_objects().pred : nullptr, + use_tdt ? &decoder_objects().joint : nullptr, + enc_outs[b], d_model, valid_Tout[b], use_tdt); } return outs; } @@ -391,7 +548,8 @@ std::vector Model::transcribe_pcm_batch( // Transcription (text + per-word timestamps + tokens). Mirrors the decode tail // of transcribe_16k_with_timestamps exactly. static Transcription decode_enc_out_with_timestamps( - const ModelLoader& loader, const std::vector& enc_out, + const ModelLoader& loader, const PredictionNet* dpred, const Joint* djoint, + const std::vector& enc_out, int d_model, int Tout, bool use_tdt, float frame_sec) { const ParakeetConfig& cfg = loader.config(); Transcription result; @@ -401,8 +559,8 @@ static Transcription decode_enc_out_with_timestamps( for (int t = 0; t < Tout; ++t) for (int c = 0; c < d_model; ++c) enc_row[(size_t)t * d_model + c] = enc_out[(size_t)c * Tout + t]; - PredictionNet pred(loader); - Joint joint(loader); + const PredictionNet& pred = *dpred; + const Joint& joint = *djoint; const int max_symbols = (int)cfg.max_symbols; if (!cfg.tdt_durations.empty()) tdt_greedy(pred, joint, enc_row, Tout, d_model, cfg.tdt_durations, @@ -450,7 +608,8 @@ Transcription Model::transcribe_16k_with_timestamps( || (decoder == Decoder::kDefault && arch_prefers_tdt(cfg.arch)); Transcription result = decode_enc_out_with_timestamps( - loader_, encoded.channels_first, encoded.d_model, encoded.frames, + loader_, use_tdt ? &decoder_objects().pred : nullptr, + use_tdt ? &decoder_objects().joint : nullptr, encoded.channels_first, encoded.d_model, encoded.frames, use_tdt, frame_sec); return result; } @@ -492,8 +651,8 @@ std::vector Model::transcribe_16k_batch_with_timestamps( std::vector> encs; std::vector Ts; batch_enc_to_row_major(enc_outs, valid_Tout, d_model, encs, Ts); - PredictionNet pred(loader_); - Joint joint(loader_); + const PredictionNet& pred = decoder_objects().pred; + const Joint& joint = decoder_objects().joint; std::vector> ids; std::vector> toks; pk::transducer_greedy_batch(pred, joint, encs, Ts, d_model, @@ -512,7 +671,8 @@ std::vector Model::transcribe_16k_batch_with_timestamps( // CTC stays per-item (not a transducer; no autoregressive decode). for (int b = 0; b < mb.B; ++b) outs[b] = decode_enc_out_with_timestamps( - loader_, enc_outs[b], d_model, valid_Tout[b], use_tdt, frame_sec); + loader_, use_tdt ? &decoder_objects().pred : nullptr, + use_tdt ? &decoder_objects().joint : nullptr, enc_outs[b], d_model, valid_Tout[b], use_tdt, frame_sec); } return outs; } @@ -545,8 +705,8 @@ std::vector Model::transcribe_16k_nbest( enc_row[(size_t)t * encoded.d_model + c] = encoded.channels_first[(size_t)c * encoded.frames + t]; - PredictionNet pred(loader_); - Joint joint(loader_); + const PredictionNet& pred = decoder_objects().pred; + const Joint& joint = decoder_objects().joint; std::vector beam = tdt_beam_search( pred, joint, enc_row, encoded.frames, encoded.d_model, cfg.tdt_durations, (int)cfg.blank_id, diff --git a/src/model.hpp b/src/model.hpp index e3d280e..14005ff 100644 --- a/src/model.hpp +++ b/src/model.hpp @@ -1,10 +1,15 @@ #pragma once #include "parakeet.h" // pk::Decoder +#include "joint.hpp" #include "model_loader.hpp" +#include "prediction.hpp" #include "tdt.hpp" // pk::TdtBeamToken #include "transcription.hpp" // pk::Transcription +#include "vad_head.hpp" +#include "vad_segmenter.hpp" #include +#include #include #include @@ -26,7 +31,10 @@ struct NBestTranscription { // // The component objects (MelFrontend, Encoder, PredictionNet, Joint, ...) are // lightweight views over the ModelLoader (they hold `const ModelLoader&`), so -// they are constructed per call; the expensive part — parsing the GGUF and +// they are constructed per call, except the transducer decoder objects +// (PredictionNet, Joint), which are built once on first use and shared by every +// decode (they are immutable after construction; see decoder_objects()). The +// expensive part — parsing the GGUF and // mapping every weight tensor — happens exactly once, in load(). class Model { public: @@ -120,6 +128,23 @@ class Model { const std::string& target_lang = "") const; const ParakeetConfig& config() const { return loader_.config(); } + // Per-frame speech probability from the model's own VAD head (80 ms frames, + // cfg.vad.frame_sec). `v` overrides the head wiring for experiments; nullptr + // uses the defaults. Throws std::runtime_error if the model has no VAD head. + std::vector vad_probabilities(const std::vector& pcm16k, + const VadVariant* v = nullptr) const; + + // Transcribe long audio in VAD-cut segments (see vad_segmenter.hpp). Audio no + // longer than opts.max_seg_sec takes the plain path. Requires a model with a + // VAD head (throws std::runtime_error("model has no VAD head") otherwise). + std::string transcribe_pcm_vad(const std::vector& pcm, int sample_rate, + Decoder decoder = Decoder::kDefault, + const std::string& target_lang = "", + const SegmenterOpts& opts = SegmenterOpts()) const; + Transcription transcribe_pcm_vad_with_timestamps( + const std::vector& pcm, int sample_rate, + Decoder decoder = Decoder::kDefault, const std::string& target_lang = "", + const SegmenterOpts& opts = SegmenterOpts()) const; // The underlying loaded GGUF. Exposed so the streaming C-API can build a // pk::StreamingSession (and a MelFrontend) over the same load-once model. @@ -163,6 +188,21 @@ class Model { const std::string& target_lang = "") const; ModelLoader loader_; + + // The transducer decoder objects, built once per model on first use + // (thread-safe) so the 21 MB embedding table is not re-copied per + // utterance. PredictionNet and Joint are read-only after construction + // (the embedding table is filled once under a std::once_flag), so + // concurrent decodes may share them. Only valid for transducer models. + struct DecoderObjects { + PredictionNet pred; + Joint joint; + explicit DecoderObjects(const ModelLoader& ml) : pred(ml), joint(ml) {} + }; + const DecoderObjects& decoder_objects() const; + + mutable std::once_flag decoder_once_; + mutable std::unique_ptr decoder_; }; } // namespace pk diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ae33991..96004f4 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -5,6 +5,7 @@ #include "ggml-alloc.h" #include "ggml-cpu.h" #include "gguf.h" +#include #include #include #include @@ -170,6 +171,29 @@ bool ModelLoader::load(const std::string& path){ cfg_.prompt.dict_keys = kv_str_arr(gguf_, "parakeet.prompt.dictionary.keys"); cfg_.prompt.dict_vals = kv_i32_arr(gguf_, "parakeet.prompt.dictionary.values"); } + cfg_.ternary.present = kv_bool(gguf_, "parakeet.ternary.present", false); + if(cfg_.ternary.present){ + cfg_.ternary.group_size = kv_u32(gguf_, "parakeet.ternary.group_size", 128); + if(cfg_.ternary.group_size != 128){ + PK_LOG("invalid packed ternary GGUF: parakeet.ternary.group_size is %u, only 128 is supported", + (unsigned)cfg_.ternary.group_size); + return false; + } + } + cfg_.vad.present = kv_bool(gguf_, "parakeet.vad.present", false); + if(cfg_.vad.present){ + cfg_.vad.d_in = kv_u32(gguf_, "parakeet.vad.d_in", 0); + cfg_.vad.hidden = kv_u32(gguf_, "parakeet.vad.hidden", 0); + cfg_.vad.kernel = kv_u32(gguf_, "parakeet.vad.kernel", 0); + cfg_.vad.frame_sec = kv_f32(gguf_, "parakeet.vad.frame_sec", 0.08f); + if(!std::isfinite(cfg_.vad.frame_sec) || !(cfg_.vad.frame_sec > 0.0f) || + cfg_.vad.d_in==0 || cfg_.vad.hidden==0 || cfg_.vad.kernel==0 || (cfg_.vad.kernel % 2)==0){ + PK_LOG("invalid VAD config: frame_sec=%g d_in=%u hidden=%u kernel=%u (need finite frame_sec > 0, " + "non-zero sizes and an odd kernel)", (double)cfg_.vad.frame_sec, + (unsigned)cfg_.vad.d_in, (unsigned)cfg_.vad.hidden, (unsigned)cfg_.vad.kernel); + return false; + } + } if(cfg_.att_context_style != "regular"){ StreamingCfg& s = cfg_.streaming; s.chunk_size = kv_i32_arr(gguf_, "parakeet.streaming.chunk_size"); @@ -186,6 +210,15 @@ bool ModelLoader::load(const std::string& path){ cfg_.n_fft = kv_u32(gguf_, "parakeet.preprocessor.n_fft"); cfg_.win_length = kv_u32(gguf_, "parakeet.preprocessor.win_length"); cfg_.hop_length = kv_u32(gguf_, "parakeet.preprocessor.hop_length"); + if(cfg_.vad.present){ + const double enc_frame = (double)cfg_.hop_length * (double)cfg_.subsampling_factor / (double)cfg_.sample_rate; + const double ratio = enc_frame > 0.0 ? (double)cfg_.vad.frame_sec / enc_frame : 0.0; + if(!(enc_frame > 0.0) || std::fabs(ratio - std::round(ratio)) > 1e-3 || std::round(ratio) < 1.0){ + PK_LOG("invalid VAD config: frame_sec=%g is not a whole multiple of the encoder frame (%g s)", + (double)cfg_.vad.frame_sec, enc_frame); + return false; + } + } cfg_.preemph = kv_f32(gguf_, "parakeet.preprocessor.preemph", 0.0f); cfg_.mag_power = kv_f32(gguf_, "parakeet.preprocessor.mag_power", 2.0f); cfg_.normalize = kv_str(gguf_, "parakeet.preprocessor.normalize", "per_feature"); diff --git a/src/model_loader.hpp b/src/model_loader.hpp index 7cfedae..121a4ef 100644 --- a/src/model_loader.hpp +++ b/src/model_loader.hpp @@ -1,4 +1,5 @@ #pragma once +#include #include #include #include @@ -44,6 +45,20 @@ struct PromptCfg { // paths so both reject typos identically (matches the C-API contract). int resolve_index_or_throw(const std::string& target_lang) const; }; +// Packed ternary encoder linears (moondream/parakeet-redux). present=false for +// every other model; the encoder then uses ordinary .weight tensors. +struct TernaryCfg { + bool present = false; + uint32_t group_size = 128; // scales are per (row, 128 columns) +}; +// Voice-activity head on the subsampler output (moondream ultra/redux). +struct VadCfg { + bool present = false; + uint32_t d_in = 0; // input channels (subsampler d_model) + uint32_t hidden = 0; + uint32_t kernel = 0; + float frame_sec = 0.08f; // seconds per head output frame +}; struct ParakeetConfig { std::string arch; // encoder @@ -59,6 +74,8 @@ struct ParakeetConfig { bool use_bias=true; // false for nemotron (encoder linears have no bias) StreamingCfg streaming; PromptCfg prompt; // prompt conditioning (present=false for non-prompt) + TernaryCfg ternary; // present=false unless parakeet.ternary.present + VadCfg vad; // present=false unless parakeet.vad.present // preprocessor uint32_t sample_rate=16000, n_mels=0, n_fft=0, win_length=0, hop_length=0; float preemph=0.0f, mag_power=2.0f, log_zero_guard=0.0f; @@ -115,6 +132,14 @@ class ModelLoader { const std::vector& tokenizer_pieces() const { return cfg_.tokenizer_pieces; } ggml_tensor* tensor(const std::string& name) const; // nullptr if absent ggml_context* ggml_ctx() const { return ctx_; } + // True iff any tensor name ends with suffix (a scan of every tensor). + bool has_tensor_with_suffix(const std::string& suffix) const { + for(const auto& kv : tensors_){ + const std::string& n = kv.first; + if(n.size() >= suffix.size() && n.compare(n.size()-suffix.size(), suffix.size(), suffix)==0) return true; + } + return false; + } // Give every weight tensor a CPU backend buffer (ONCE), so graphs can // reference the loader's tensors DIRECTLY as leaves with zero per-call @@ -127,7 +152,12 @@ class ModelLoader { // the same CPU backend the compute path uses. Returns false on failure. bool realize_weights(ggml_backend_t backend); bool weights_realized() const { return weights_buf_ != nullptr; } + // Opaque per-loader slot for the ternary weight cache (see ternary.cpp). + // Owned here so cached repacked weights die with the loader that owns the + // tensors they were built from. + std::shared_ptr& ternary_store() const { return ternary_store_; } private: + mutable std::shared_ptr ternary_store_; ParakeetConfig cfg_; gguf_context* gguf_ = nullptr; ggml_context* ctx_ = nullptr; diff --git a/src/parakeet_capi.cpp b/src/parakeet_capi.cpp index 8998f4b..7c2d1da 100644 --- a/src/parakeet_capi.cpp +++ b/src/parakeet_capi.cpp @@ -13,6 +13,7 @@ #include "transcription.hpp" // pk::Transcription, pk::Word #include "transcription_json.hpp" +#include "audio_io.hpp" #include #include @@ -444,6 +445,42 @@ extern "C" char* parakeet_capi_transcribe_path_json(parakeet_ctx* ctx, } } +extern "C" char* parakeet_capi_transcribe_path_json_vad(parakeet_ctx* ctx, + const char* wav_path, + int decoder) { + if (!ctx) return nullptr; + if (!ctx->model) { + ctx->last_error = ctx->diar + ? "context holds a diarization model; use parakeet_capi_diarize_*" + : "context has no loaded model"; + return nullptr; + } + if (!wav_path) { ctx->last_error = "wav_path is NULL"; return nullptr; } + try { + pk::Audio audio; + if (!pk::load_audio_16k_mono(wav_path, audio)) { + ctx->last_error = std::string("failed to load audio: ") + wav_path; + return nullptr; + } + pk::Transcription tr = ctx->model->transcribe_pcm_vad_with_timestamps( + audio.samples, audio.sample_rate, to_decoder(decoder)); + 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::transcription_to_json(tr, 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_batch_json_lang(parakeet_ctx* ctx, const float* samples_concat, const int* n_samples, int n_clips, int sample_rate, int decoder, const char* target_lang) { diff --git a/src/prediction.cpp b/src/prediction.cpp index 3ef8066..6faf79a 100644 --- a/src/prediction.cpp +++ b/src/prediction.cpp @@ -28,6 +28,23 @@ PredState PredictionNet::zero_state() const { return s; } +// Fetch the embedding table to the host once (device-safe), thread-safe so a +// PredictionNet can be shared by concurrent decodes. Ensure the loader's +// weights have a backend buffer first (idempotent) so the tensor is readable +// via ggml_backend_tensor_get even when step()/forward() runs before the +// encoder graph has realized the weights. After the call_once the table is +// read-only. +const float* PredictionNet::embed_table() const { + std::call_once(embed_once_, [this]() { + pk::ensure_weights_realized(ml_); + ggml_tensor* emb = ml_.tensor("decoder.prediction.embed.weight"); + assert(emb && "missing decoder.prediction.embed.weight"); + embed_host_.resize((size_t)vocab_p1_ * H_); + ggml_backend_tensor_get(emb, embed_host_.data(), 0, ggml_nbytes(emb)); + }); + return embed_host_.data(); +} + // --------------------------------------------------------------------------- // Advance the stacked LSTM by one token, as a single ggml graph that runs on // whatever backend pk::Backend selected (CPU or device). The embedding table is @@ -48,23 +65,13 @@ void PredictionNet::step(int32_t token_id, bool is_sos, const int H = H_; const int L = n_layers_; - // Lazily fetch the embedding table to the host (device-safe). Ensure the - // loader's weights have a backend buffer first (idempotent) so the tensor - // is readable via ggml_backend_tensor_get even when step()/forward() is - // exercised before the encoder graph has realized the weights. - if (embed_host_.empty()) { - pk::ensure_weights_realized(ml_); - ggml_tensor* emb = ml_.tensor("decoder.prediction.embed.weight"); - assert(emb && "missing decoder.prediction.embed.weight"); - embed_host_.resize((size_t)vocab_p1_ * H); - ggml_backend_tensor_get(emb, embed_host_.data(), 0, ggml_nbytes(emb)); - } + const float* embed_host = embed_table(); // Layer-0 input: zeros for SOS, else the embedding row for token_id. std::vector x0((size_t)H, 0.0f); if (!is_sos) { assert(token_id >= 0 && token_id < vocab_p1_ && "embedding id out of range"); - std::memcpy(x0.data(), &embed_host_[(size_t)token_id * H], + std::memcpy(x0.data(), embed_host + (size_t)token_id * H, (size_t)H * sizeof(float)); } @@ -128,22 +135,14 @@ void PredictionNet::step_batch(const std::vector& token_ids, const int N = (int)token_ids.size(); assert(N > 0 && (int)is_sos.size() == N && "batch size mismatch"); - // Lazily fetch the embedding table to the host (device-safe), exactly as - // step() does. - if (embed_host_.empty()) { - pk::ensure_weights_realized(ml_); - ggml_tensor* emb = ml_.tensor("decoder.prediction.embed.weight"); - assert(emb && "missing decoder.prediction.embed.weight"); - embed_host_.resize((size_t)vocab_p1_ * H); - ggml_backend_tensor_get(emb, embed_host_.data(), 0, ggml_nbytes(emb)); - } + const float* embed_host = embed_table(); // Layer-0 input [H*N]: zeros for SOS items, else the embedding row. std::vector x0((size_t)H * N, 0.0f); for (int n = 0; n < N; ++n) { if (!is_sos[n]) { assert(token_ids[n] >= 0 && token_ids[n] < vocab_p1_ && "embedding id out of range"); - std::memcpy(&x0[(size_t)n * H], &embed_host_[(size_t)token_ids[n] * H], + std::memcpy(&x0[(size_t)n * H], embed_host + (size_t)token_ids[n] * H, (size_t)H * sizeof(float)); } } diff --git a/src/prediction.hpp b/src/prediction.hpp index 6e5b8ed..816939c 100644 --- a/src/prediction.hpp +++ b/src/prediction.hpp @@ -1,4 +1,5 @@ #pragma once +#include #include "model_loader.hpp" #include #include @@ -96,9 +97,15 @@ class PredictionNet { int vocab_p1_; // vocab + 1 (embedding rows) int n_layers_; // pred_rnn_layers (stacked LSTM layers) + // Returns the host-side embedding table, fetching it on first use. + const float* embed_table() const; + // Host-side copy of the embedding table, lazily fetched on the first step() // via ggml_backend_tensor_get (works for both CPU and device-resident - // weights). [vocab_p1_ * H_], row-major: embed_host_[id*H_ + h]. + // weights). [vocab_p1_ * H_], row-major: embed_host_[id*H_ + h]. Filled + // exactly once under embed_once_ and read-only afterwards, so a shared + // PredictionNet is safe to use from several threads. + mutable std::once_flag embed_once_; mutable std::vector embed_host_; }; diff --git a/src/relpos_attention.cpp b/src/relpos_attention.cpp index c82038e..585acd6 100644 --- a/src/relpos_attention.cpp +++ b/src/relpos_attention.cpp @@ -1,11 +1,13 @@ #include "relpos_attention.hpp" #include "ggml_graph.hpp" #include "backend.hpp" +#include "ternary.hpp" #include "ggml.h" #include #include #include #include +#include #include namespace pk { @@ -21,6 +23,31 @@ static ggml_tensor* clone_weight(ggml_context* ctx, const ModelLoader& ml, return pk::clone_weight(ctx, ml, name.c_str()); } +// One projection for every attention path (single, batched, local, chunked). +// Packed ternary GGUFs carry `.qweight` + `.scales` and no +// `.weight`; everything else goes through ggml_mul_mat on the zero-copy +// loader tensor. The bias is added only when requested AND present. A weight +// that is neither packed nor plain is a hard error, never a null into ggml. +static ggml_tensor* attn_linear(ggml_context* ctx, const ModelLoader& ml, + const std::string& pre, const char* w, const char* b, + ggml_tensor* in) { + std::string base = pre + w; // e.g. "...self_attn.linear_q.weight" + base.resize(base.size() - 7); // drop ".weight" + ggml_tensor* y; + if (has_ternary(ml, base)) { + y = ternary_linear(ctx, ml, base, in); // [out, *] + } else { + if (!ml.tensor(pre + w)) + throw std::runtime_error("missing encoder weight " + pre + w); + y = ggml_mul_mat(ctx, clone_weight(ctx, ml, pre + w), in); // [out, *] + } + if (b && ml.tensor(pre + b)) { + ggml_tensor* B = clone_weight(ctx, ml, pre + b); + y = ggml_add(ctx, y, B); // broadcast [out] over cols + } + return y; +} + RelPosAttention::RelPosAttention(const ModelLoader& ml, int layer_idx) : ml_(ml), layer_idx_(layer_idx) { d_model_ = (int)ml.config().d_model; @@ -59,13 +86,7 @@ ggml_tensor* RelPosAttention::build_graph(ggml_context* ctx, ggml_tensor* xt, // attention linears with bias=False in some checkpoints // (parakeet-tdt-0.6b-v2/-v3) and bias=True in others (110m). auto linear = [&](const char* w, const char* b, ggml_tensor* in) { - ggml_tensor* W = clone_weight(ctx, ml, pre + w); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); // [out, *] - if (b && ml.tensor(pre + b)) { - ggml_tensor* B = clone_weight(ctx, ml, pre + b); - y = ggml_add(ctx, y, B); // broadcast [out] over cols - } - return y; + return attn_linear(ctx, ml, pre, w, b, in); }; ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); // [D, T] ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); // [D, T] @@ -194,13 +215,7 @@ ggml_tensor* RelPosAttention::build_graph_batched( // attention linears with bias=False in some checkpoints // (parakeet-tdt-0.6b-v2/-v3) and bias=True in others (110m). auto linear = [&](const char* w, const char* b, ggml_tensor* in) { - ggml_tensor* W = clone_weight(ctx, ml, pre + w); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); // [out, *] - if (b && ml.tensor(pre + b)) { - ggml_tensor* B = clone_weight(ctx, ml, pre + b); - y = ggml_add(ctx, y, B); // broadcast [out] over cols - } - return y; + return attn_linear(ctx, ml, pre, w, b, in); }; // xt is [D, T, B]; mul_mat batches over ne2 -> q/k/v are [D, T, B]. pe is // shared [D, P] (NO batch) -> p is [D, P]. @@ -347,10 +362,7 @@ ggml_tensor* RelPosAttention::build_graph_batched_local( const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn."; const ModelLoader& ml = ml_; auto linear = [&](const char* wn, const char* bn, ggml_tensor* in) { - ggml_tensor* W = clone_weight(ctx, ml, pre + wn); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); - if (bn && ml.tensor(pre + bn)) y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + bn)); - return y; + return attn_linear(ctx, ml, pre, wn, bn, in); }; ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); // [D, T, B] ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); @@ -480,10 +492,7 @@ ggml_tensor* RelPosAttention::build_graph_local(ggml_context* ctx, ggml_tensor* const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn."; const ModelLoader& ml = ml_; auto linear = [&](const char* wn, const char* bn, ggml_tensor* in) { - ggml_tensor* W = clone_weight(ctx, ml, pre + wn); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); - if (bn && ml.tensor(pre + bn)) y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + bn)); - return y; + return attn_linear(ctx, ml, pre, wn, bn, in); }; ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); @@ -560,10 +569,7 @@ ggml_tensor* RelPosAttention::build_graph_local(ggml_context* ctx, ggml_tensor* qm.data(), qm.size() * sizeof(float)); merged = ggml_mul(ctx, merged, qmask); } - ggml_tensor* Wo = clone_weight(ctx, ml, pre + "linear_out.weight"); - ggml_tensor* y = ggml_mul_mat(ctx, Wo, merged); - if (ml.tensor(pre + "linear_out.bias")) - y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + "linear_out.bias")); + ggml_tensor* y = attn_linear(ctx, ml, pre, "linear_out.weight", "linear_out.bias", merged); return y; // [D, T] } } @@ -618,10 +624,7 @@ ggml_tensor* RelPosAttention::build_graph_local_chunked( const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn."; const ModelLoader& ml = ml_; auto linear = [&](const char* wn, const char* bn, ggml_tensor* in) { - ggml_tensor* W = clone_weight(ctx, ml, pre + wn); - ggml_tensor* y = ggml_mul_mat(ctx, W, in); - if (bn && ml.tensor(pre + bn)) y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + bn)); - return y; + return attn_linear(ctx, ml, pre, wn, bn, in); }; ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); @@ -725,10 +728,7 @@ ggml_tensor* RelPosAttention::build_graph_local_chunked( qm.data(), qm.size() * sizeof(float)); merged = ggml_mul(ctx, merged, qmask); } - ggml_tensor* Wo = clone_weight(ctx, ml, pre + "linear_out.weight"); - ggml_tensor* y = ggml_mul_mat(ctx, Wo, merged); - if (ml.tensor(pre + "linear_out.bias")) - y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + "linear_out.bias")); + ggml_tensor* y = attn_linear(ctx, ml, pre, "linear_out.weight", "linear_out.bias", merged); return y; // [D, T] } diff --git a/src/streaming_encoder.cpp b/src/streaming_encoder.cpp index 44f289d..8254561 100644 --- a/src/streaming_encoder.cpp +++ b/src/streaming_encoder.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -28,6 +29,8 @@ static ggml_tensor* clone_weight_opt(ggml_context* ctx, const ModelLoader& ml, } StreamingEncoder::StreamingEncoder(const ModelLoader& ml) : ml_(ml) { + if (ml.config().ternary.present || ml.has_tensor_with_suffix(".qweight")) + throw std::runtime_error("streaming does not support packed ternary GGUFs; convert with --ternary dequant"); const ParakeetConfig& c = ml.config(); d_model_ = (int)c.d_model; n_layers_ = (int)c.n_layers; diff --git a/src/subsampling.cpp b/src/subsampling.cpp index b3272d3..89ce169 100644 --- a/src/subsampling.cpp +++ b/src/subsampling.cpp @@ -69,8 +69,10 @@ ggml_tensor* Subsampling::build_graph_batched(ggml_context* ctx, // the leading ggml_pad_ext (lp1=2/rp1=1 on time) uniformly across the batch, // and the per-item trailing-pad time masking (mask_time on the batch axis) // plus the all_paddings=3 valid-length recurrence reproduce, per item, the - // exact standalone causal boundary. A clip in a B>1 batch is byte-identical - // to the same clip transcribed standalone (see test_subsampling_batch_causal). + // exact standalone causal boundary. A clip in a B>1 batch gets the same + // transcript as the same clip transcribed standalone on the test clips (see + // test_subsampling_batch_causal); float values are close but not guaranteed + // bit-identical, since ggml picks different matmul kernels for B>1. // --- Input (host-side): ggml conv data layout is [W=feat, H=T, IC=1, N=B]. // NeMo conv input is [B,1,T,feat] (H=T, W=feat). We must feed diff --git a/src/ternary.cpp b/src/ternary.cpp new file mode 100644 index 0000000..8804e65 --- /dev/null +++ b/src/ternary.cpp @@ -0,0 +1,405 @@ +#include "ternary.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ggml.h" +#include "common.hpp" +#include "model_loader.hpp" +#include "ternary_kernels.hpp" + +namespace pk { + +namespace { + +// For every byte value 0..255 its five base-3 digits, least significant first, +// computed exactly like the scalar reference (v % 3, v /= 3), so the invalid +// bytes 243..255 give the same digits as before. +struct TritTable { + uint8_t d[256][5]; + TritTable() { + for (int b = 0; b < 256; ++b) { + int v = b; + for (int i = 0; i < 5; ++i) { + d[b][i] = (uint8_t)(v % 3); + v /= 3; + } + } + } +}; + +const TritTable& trit_table() { + static const TritTable t; // thread safe initialization + return t; +} + +} // namespace + +void ternary_repack(const uint8_t* q, const uint16_t* s, int N, int K, TernaryWeight& out) { + GGML_ASSERT(K > 0 && K % kTernaryGroup == 0); + const int G = K / kTernaryGroup; + const int nb = (K + 4) / 5; + out.N = N; + out.K = K; + const int B = out.row_blocks(); + out.planes.assign((size_t)B * G * 512, 0); + out.scales.assign((size_t)B * G * kTernaryRowBlock, 0.0f); + const TritTable& tab = trit_table(); + // 5 spare bytes: the last upstream byte may hold digits past K, written and ignored + std::vector row((size_t)K + 8); + for (int n = 0; n < N; ++n) { + const int b = n / kTernaryRowBlock, i = n % kTernaryRowBlock; + for (int g = 0; g < G; ++g) + out.scales[((size_t)b * G + g) * kTernaryRowBlock + i] = ggml_fp16_to_fp32(s[(size_t)n * G + g]); + const uint8_t* qr = q + (size_t)n * nb; + for (int bb = 0; bb < nb; ++bb) std::memcpy(&row[(size_t)bb * 5], tab.d[qr[bb]], 5); + for (int g = 0; g < G; ++g) { + const uint8_t* src = row.data() + (size_t)g * kTernaryGroup; + uint8_t* dst = out.planes.data() + ((size_t)b * G + g) * 512 + 4 * i; + // step st, byte j: the four codes at 16*st + 4*p + j, p = 0..3, in bit pairs + for (int st = 0; st < 8; ++st) { + const uint8_t* e = src + st * 16; + for (int j = 0; j < 4; ++j) + dst[st * 64 + j] = (uint8_t)(e[j] | (e[4 + j] << 2) | (e[8 + j] << 4) | (e[12 + j] << 6)); + } + } + } +} + +void ternary_dequant(const uint8_t* q, const uint16_t* s, int N, int K, std::vector& W) { + const int G = (K + kTernaryGroup - 1) / kTernaryGroup; + const int nb = (K + 4) / 5; + W.assign((size_t)N * K, 0.0f); + for (int n = 0; n < N; ++n) + for (int b = 0; b < nb; ++b) { + int v = q[(size_t)n * nb + b]; + for (int d = 0; d < 5; ++d) { + const int i = b * 5 + d; + if (i < K) + W[(size_t)n * K + i] = + ggml_fp16_to_fp32(s[(size_t)n * G + i / kTernaryGroup]) * (float)(v % 3 - 1); + v /= 3; + } + } +} + +void ternary_quant_rows_ref(const float* x, int K, int t0, int t1, uint8_t* act) { + const int G = K / kTernaryGroup; + const size_t rb = ternary_act_row_bytes(K); + for (int t = t0; t < t1; ++t) { + const float* xr = x + (size_t)t * K; + uint8_t* row = act + (size_t)t * rb; + int8_t* q = reinterpret_cast(row); + float* sc = reinterpret_cast(row + K); + int32_t* gs = reinterpret_cast(row + K + 4); + float amax = 0.0f; + for (int k = 0; k < K; ++k) amax = std::max(amax, std::fabs(xr[k])); + const float s = amax > 0.0f ? amax / 127.0f : 0.0f; + const float inv = s > 0.0f ? 1.0f / s : 0.0f; + *sc = s; + for (int g = 0; g < G; ++g) { + int32_t sum = 0; + for (int j = 0; j < kTernaryGroup; ++j) { + // clamp before rounding (same result, since +-127 are + // integers), which also maps a NaN input to -127 + const float f = std::min(127.0f, std::max(-127.0f, xr[g * kTernaryGroup + j] * inv)); + const int v = (int)std::lrintf(f); + q[g * kTernaryGroup + j] = (int8_t)v; + sum += v; + } + gs[g] = sum; + } + } +} + +void ternary_matmul_rows_ref(const TernaryWeight& w, const uint8_t* act, int T, float* y, + int r0, int r1) { + const int N = w.N, K = w.K, G = w.groups(); + const size_t rb = ternary_act_row_bytes(K); + for (int n = r0; n < r1; ++n) + for (int t = 0; t < T; ++t) { + const uint8_t* row = act + (size_t)t * rb; + const int8_t* q = reinterpret_cast(row); + const float sa = *reinterpret_cast(row + K); + const int32_t* gs = reinterpret_cast(row + K + 4); + float acc = 0.0f; + for (int g = 0; g < G; ++g) { + int32_t s = 0; + for (int j = 0; j < kTernaryGroup; ++j) + s += w.code(n, g * kTernaryGroup + j) * (int32_t)q[g * kTernaryGroup + j]; + acc += w.scale(n, g) * (float)(s - gs[g]); + } + y[(size_t)t * N + n] = acc * sa; + } +} + +namespace { + +const TernaryKernel kScalarKernel{"scalar", ternary_matmul_rows_ref}; + +const TernaryKernel* pick_kernel() { + static const TernaryKernel* chosen = [] { + if (const char* e = std::getenv("PARAKEET_TERNARY_KERNEL")) { + const TernaryKernel* want = nullptr; + if (!std::strcmp(e, "scalar")) want = &kScalarKernel; + else if (!std::strcmp(e, "avx2")) want = ternary_kernel_x86_avx2(); + else if (!std::strcmp(e, "vnni")) want = ternary_kernel_x86_vnni(); + else if (!std::strcmp(e, "neon")) want = ternary_kernel_neon(); + if (want) return want; + PK_LOG("PARAKEET_TERNARY_KERNEL=%s is unavailable here; using automatic selection", e); + } + if (const TernaryKernel* k = ternary_kernel_x86_vnni()) return k; + if (const TernaryKernel* k = ternary_kernel_neon()) return k; + if (const TernaryKernel* k = ternary_kernel_x86_avx2()) return k; + return &kScalarKernel; + }(); + return chosen; +} + +} // namespace + +std::vector ternary_all_kernels() { + std::vector v{&kScalarKernel}; + for (const TernaryKernel* k : {ternary_kernel_x86_avx2(), ternary_kernel_x86_vnni(), ternary_kernel_neon()}) + if (k) v.push_back(k); + return v; +} + +std::vector ternary_all_quants() { + static const TernaryQuant ref{"scalar", ternary_quant_rows_ref}; + std::vector v{&ref}; + for (const TernaryQuant* q : {ternary_quant_x86_avx2(), ternary_quant_x86_avx512()}) + if (q) v.push_back(q); + return v; +} + +void ternary_quant_rows(const float* x, int K, int t0, int t1, uint8_t* act) { + static const TernaryQuantFn fn = ternary_all_quants().back()->fn; + fn(x, K, t0, t1, act); +} + +void ternary_matmul_rows(const TernaryWeight& w, const uint8_t* act, int T, float* y, int r0, int r1) { + pick_kernel()->fn(w, act, T, y, r0, r1); +} + +const char* ternary_kernel_name() { return pick_kernel()->name; } + +namespace { + +struct TernaryStore { + std::mutex mu; + std::unordered_map> weights; +}; + +// Validates the packed tensors of one linear before any read of their data. +// expN / expK are the shapes the model config demands (or -1 to skip). +void validate_ternary_tensors(const ggml_tensor* q, const ggml_tensor* s, const std::string& base, + int64_t expN, int64_t expK) { + const std::string who = "ternary: " + base + ": "; + if (q->type != GGML_TYPE_I8) throw std::runtime_error(who + "qweight is not I8"); + if (s->type != GGML_TYPE_F16) throw std::runtime_error(who + "scales is not F16"); + if (q->ne[2] != 1 || q->ne[3] != 1 || s->ne[2] != 1 || s->ne[3] != 1) + throw std::runtime_error(who + "tensors must be 2-D"); + const int64_t kMax = (int64_t)1 << 24; // far above any real layer, keeps all products in range + const int64_t N = q->ne[1]; + const int64_t G = s->ne[0]; + if (N <= 0 || G <= 0 || N > kMax || G > kMax) + throw std::runtime_error(who + "invalid dimensions"); + const int64_t K = G * kTernaryGroup; + if (K % kTernaryGroup != 0) throw std::runtime_error(who + "K is not a multiple of 128"); + if (s->ne[1] != N) throw std::runtime_error(who + "scales rows differ from qweight rows"); + const int64_t nb = (K + 4) / 5; + if (q->ne[0] != nb) throw std::runtime_error(who + "qweight row length does not match K"); + if ((uint64_t)ggml_nbytes(q) != (uint64_t)N * (uint64_t)nb) + throw std::runtime_error(who + "qweight byte size mismatch"); + if ((uint64_t)ggml_nbytes(s) != (uint64_t)N * (uint64_t)G * 2u) + throw std::runtime_error(who + "scales byte size mismatch"); + if (expN >= 0 && (N != expN || K != expK)) + throw std::runtime_error(who + "shape " + std::to_string(N) + "x" + std::to_string(K) + + " does not match the model config (" + std::to_string(expN) + "x" + + std::to_string(expK) + ")"); +} + +TernaryStore& store_of(const ModelLoader& ml) { + static std::mutex init_mu; + std::lock_guard lk(init_mu); + auto& slot = ml.ternary_store(); + if (!slot) slot = std::make_shared(); + return *static_cast(slot.get()); +} + +// Looks the two tensors up and validates them; throws on a malformed pair. +struct TernaryJob { + std::string base; + const uint8_t* q = nullptr; + const uint16_t* s = nullptr; + int N = 0, K = 0; + std::unique_ptr out; + std::exception_ptr err; + void run() { + try { + out = std::make_unique(); + ternary_repack(q, s, N, K, *out); + } catch (...) { + err = std::current_exception(); + } + } +}; + +TernaryJob make_job(const ModelLoader& ml, const std::string& base, int64_t expN, int64_t expK) { + const ggml_tensor* q = ml.tensor(base + ".qweight"); + const ggml_tensor* s = ml.tensor(base + ".scales"); + if (!q || !s) throw std::runtime_error("ternary: missing " + base + ".qweight/.scales"); + validate_ternary_tensors(q, s, base, expN, expK); + TernaryJob j; + j.base = base; + j.q = static_cast(q->data); + j.s = static_cast(s->data); + j.N = (int)q->ne[1]; + j.K = (int)s->ne[0] * kTernaryGroup; + return j; +} + +const TernaryWeight& weight_for(const ModelLoader& ml, const std::string& base, + int64_t expN = -1, int64_t expK = -1) { + TernaryStore& store = store_of(ml); + std::lock_guard lk(store.mu); + auto it = store.weights.find(base); + if (it != store.weights.end()) return *it->second; + TernaryJob j = make_job(ml, base, expN, expK); + j.run(); + if (j.err) std::rethrow_exception(j.err); + return *store.weights.emplace(base, std::move(j.out)).first->second; +} + +// Worker count for the load-time repack: PARAKEET_REPACK_THREADS, else +// min(hardware_concurrency, 8). 1 selects the serial path. +int repack_threads() { + if (const char* e = std::getenv("PARAKEET_REPACK_THREADS")) { + char* end = nullptr; + const long v = std::strtol(e, &end, 10); + if (end != e && v >= 1) return (int)std::min(v, 256); + } + const unsigned hc = std::thread::hardware_concurrency(); + return (int)std::min(hc ? hc : 1u, 8u); +} + +// Op 1: per-row int8 quantization of the activations. dst is I8, ne0 = row bytes. +void op_quant(ggml_tensor* dst, int ith, int nth, void*) { + const ggml_tensor* x = dst->src[0]; + const int K = (int)x->ne[0]; + const int T = (int)ggml_nrows(x); + const int t0 = (int)((int64_t)T * ith / nth); + const int t1 = (int)((int64_t)T * (ith + 1) / nth); + ternary_quant_rows(static_cast(x->data), K, t0, t1, static_cast(dst->data)); +} + +// Op 2: ternary matmul. src0 is the quantized activation tensor from op 1, +// userdata is the TernaryWeight. Threads split the output rows. +void op_matmul(ggml_tensor* dst, int ith, int nth, void* ud) { + const TernaryWeight& w = *static_cast(ud); + const int T = (int)ggml_nrows(dst); + const int r0 = (int)((int64_t)w.N * ith / nth); + const int r1 = (int)((int64_t)w.N * (ith + 1) / nth); + ternary_matmul_rows(w, static_cast(dst->src[0]->data), T, + static_cast(dst->data), r0, r1); +} + +} // namespace + +std::string ternary_flag_consistency_error(const ModelLoader& ml) { + const bool packed = ml.has_tensor_with_suffix(".qweight"); + const bool flag = ml.config().ternary.present; + if (packed && !flag) + return "this GGUF holds packed ternary tensors (.qweight) but parakeet.ternary.present is not set"; + if (flag && !packed) + return "parakeet.ternary.present is set but the GGUF holds no packed ternary tensors (.qweight)"; + return ""; +} + +bool has_ternary(const ModelLoader& ml, const std::string& base) { + return ml.tensor(base + ".qweight") != nullptr; +} + +ggml_tensor* ternary_linear(ggml_context* ctx, const ModelLoader& ml, const std::string& base, + ggml_tensor* x) { + const TernaryWeight& w = weight_for(ml, base); + if (x->type != GGML_TYPE_F32 || (int)x->ne[0] != w.K) + throw std::runtime_error("ternary: input shape mismatch for " + base); + if (!ggml_is_contiguous(x)) x = ggml_cont(ctx, x); + const int64_t T = ggml_nrows(x); + ggml_tensor* a1[1] = {x}; + ggml_tensor* act = ggml_custom_4d(ctx, GGML_TYPE_I8, (int64_t)ternary_act_row_bytes(w.K), T, 1, 1, + a1, 1, op_quant, GGML_N_TASKS_MAX, nullptr); + ggml_tensor* a2[1] = {act}; + return ggml_custom_4d(ctx, GGML_TYPE_F32, w.N, x->ne[1], x->ne[2], x->ne[3], a2, 1, op_matmul, + GGML_N_TASKS_MAX, const_cast(&w)); +} + +void ternary_prepare(const ModelLoader& ml) { + const ParakeetConfig& cfg = ml.config(); + const int64_t d = cfg.d_model, ff = cfg.ff_dim; + if (cfg.ternary.group_size != (uint32_t)kTernaryGroup) + throw std::runtime_error("ternary: group_size must be 128"); + if (d <= 0 || ff <= 0) throw std::runtime_error("ternary: invalid d_model / ff_dim"); + struct Lin2 { const char* name; int64_t N, K; }; + const Lin2 kLinears[] = { + {"feed_forward1.linear1", ff, d}, {"feed_forward1.linear2", d, ff}, + {"feed_forward2.linear1", ff, d}, {"feed_forward2.linear2", d, ff}, + {"self_attn.linear_q", d, d}, {"self_attn.linear_k", d, d}, + {"self_attn.linear_v", d, d}, {"self_attn.linear_out", d, d}, + {"self_attn.linear_pos", d, d}, {"conv.pointwise_conv1", 2 * d, d}, + {"conv.pointwise_conv2", d, d}}; + const int n_layers = (int)cfg.n_layers; + // Validate everything serially and in layer order (the first error is the + // one reported), then repack the queued weights on a small thread pool. + TernaryStore& store = store_of(ml); + std::vector jobs; + for (int i = 0; i < n_layers; ++i) { + for (const Lin2& l : kLinears) { + const std::string base = "encoder.layers." + std::to_string(i) + "." + l.name; + if (has_ternary(ml, base)) { + { + std::lock_guard lk(store.mu); + if (store.weights.count(base)) continue; + } + jobs.push_back(make_job(ml, base, l.N, l.K)); + } else if (!ml.tensor(base + ".weight")) { + throw std::runtime_error("ternary: missing weight " + base); + } + } + } + std::atomic next{0}; + auto worker = [&jobs, &next] { + for (size_t k; (k = next.fetch_add(1)) < jobs.size();) jobs[k].run(); + }; + const int nthreads = (int)std::min((size_t)repack_threads(), jobs.size()); + std::vector pool; + for (int t = 1; t < nthreads; ++t) { + try { + pool.emplace_back(worker); + } catch (...) { + break; // could not start more threads: the ones running (and this one) drain the queue + } + } + worker(); + for (std::thread& th : pool) th.join(); + for (TernaryJob& j : jobs) + if (j.err) std::rethrow_exception(j.err); + std::lock_guard lk(store.mu); + for (TernaryJob& j : jobs) store.weights.emplace(j.base, std::move(j.out)); + if (std::strcmp(ternary_kernel_name(), "scalar") == 0) + PK_LOG("packed ternary weights are running on the slow scalar kernel on this CPU/build " + "(about 1 GMAC/s); re-convert with --ternary dequant for speed"); +} + +} // namespace pk diff --git a/src/ternary.hpp b/src/ternary.hpp new file mode 100644 index 0000000..fcf9dc0 --- /dev/null +++ b/src/ternary.hpp @@ -0,0 +1,92 @@ +#pragma once +// Native ternary (1.58 bit) linear layers for moondream/parakeet-redux. +// See docs/superpowers/specs/2026-09-29-hf-ternary-vad-design.md. +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +namespace pk { + +class ModelLoader; + +constexpr int kTernaryGroup = 128; +constexpr int kTernaryRowBlock = 16; + +// Repacked weight, N output rows by K input columns (K % 128 == 0). Rows are +// padded to a multiple of 16 with zero codes and zero scales, and stored in +// blocks of 16 rows so that a kernel can keep one output row per vector lane. +// planes: per row block b, per group g, per step s (0..7, 16 elements each), +// 64 bytes. Byte 4*i + j of a step holds, in bit pairs 0-1, 2-3, 4-5, 6-7 +// (plane p = 0..3), the code (0,1,2) of row 16*b + i at element +// 128*g + 16*s + 4*p + j. scales: per row block, per group, 16 floats (one per +// row of the block). w[n][k] = scale(n, k/128) * (code(n, k) - 1). +struct TernaryWeight { + int N = 0; + int K = 0; + std::vector planes; // row_blocks() * groups() * 512 bytes + std::vector scales; // row_blocks() * groups() * 16 + int groups() const { return K / kTernaryGroup; } + int row_blocks() const { return (N + kTernaryRowBlock - 1) / kTernaryRowBlock; } + // The layout above, spelled out. The scalar reference uses these. + int code(int n, int k) const { + const int b = n / kTernaryRowBlock, i = n % kTernaryRowBlock; + const int g = k / kTernaryGroup, r = k % kTernaryGroup; + const int s = r / 16, p = (r % 16) / 4, j = r % 4; + const uint8_t byte = planes[((size_t)b * groups() + g) * 512 + (size_t)s * 64 + 4 * i + j]; + return (byte >> (2 * p)) & 3; + } + float scale(int n, int g) const { + const int b = n / kTernaryRowBlock, i = n % kTernaryRowBlock; + return scales[((size_t)b * groups() + g) * kTernaryRowBlock + i]; + } +}; + +// Upstream packing (5 trits per byte, base 3, LSD first) plus F16 group scales +// to the kernel layout. qweight is N rows of ceil(K/5) bytes. +void ternary_repack(const uint8_t* qweight, const uint16_t* scales_f16, int N, int K, + TernaryWeight& out); +// Float reference dequantization, W becomes N*K row-major. +void ternary_dequant(const uint8_t* qweight, const uint16_t* scales_f16, int N, int K, + std::vector& W); + +// Quantized activations, one row per token: K int8 values, one float scale, +// K/128 int32 group sums of the int8 values. +inline size_t ternary_act_row_bytes(int K) { + return (size_t)K + 4 + 4 * (size_t)(K / kTernaryGroup); +} +// x is the full [T][K] float input; writes rows [t0, t1) of act. The reference +// defines the bytes; the dispatching entry point picks a vector version for +// this CPU that writes identical bytes. +void ternary_quant_rows_ref(const float* x, int K, int t0, int t1, uint8_t* act); +void ternary_quant_rows(const float* x, int K, int t0, int t1, uint8_t* act); + +// y[t*N + n] for n in [r0, r1) and all t < T. The reference defines correctness; +// the dispatching entry point picks the best kernel for this CPU. +void ternary_matmul_rows_ref(const TernaryWeight& w, const uint8_t* act, int T, float* y, + int r0, int r1); +void ternary_matmul_rows(const TernaryWeight& w, const uint8_t* act, int T, float* y, + int r0, int r1); +const char* ternary_kernel_name(); + +// True iff the GGUF holds .qweight (a packed ternary linear). +bool has_ternary(const ModelLoader& ml, const std::string& base); +// nn.Linear over x (f32, ne = [K, ...]) with the packed weight .qweight / +// .scales. Returns f32 with ne = [N, ...]. Bias is the caller's job. +ggml_tensor* ternary_linear(ggml_context* ctx, const ModelLoader& ml, const std::string& base, + ggml_tensor* x); + +// "" when parakeet.ternary.present agrees with the tensors: a packed tensor +// (any name ending in .qweight) exists iff the flag is set. Otherwise a message +// naming the mismatch. Scans every tensor, not just layer 0. +std::string ternary_flag_consistency_error(const ModelLoader& ml); + +// Validate and repack every packed linear of the encoder up front, so graph +// building never throws. Throws std::runtime_error on a malformed or +// half-converted file. +void ternary_prepare(const ModelLoader& ml); + +} // namespace pk diff --git a/src/ternary_kernels.hpp b/src/ternary_kernels.hpp new file mode 100644 index 0000000..7f62f37 --- /dev/null +++ b/src/ternary_kernels.hpp @@ -0,0 +1,30 @@ +#pragma once +#include + +#include "ternary.hpp" + +namespace pk { + +using TernaryRowsFn = void (*)(const TernaryWeight&, const uint8_t* act, int T, float* y, + int r0, int r1); +struct TernaryKernel { + const char* name; + TernaryRowsFn fn; +}; + +using TernaryQuantFn = void (*)(const float* x, int K, int t0, int t1, uint8_t* act); +struct TernaryQuant { + const char* name; + TernaryQuantFn fn; +}; + +const TernaryKernel* ternary_kernel_x86_vnni(); +const TernaryKernel* ternary_kernel_x86_avx2(); +const TernaryKernel* ternary_kernel_neon(); +std::vector ternary_all_kernels(); + +const TernaryQuant* ternary_quant_x86_avx512(); +const TernaryQuant* ternary_quant_x86_avx2(); +std::vector ternary_all_quants(); + +} // namespace pk diff --git a/src/ternary_kernels_neon.cpp b/src/ternary_kernels_neon.cpp new file mode 100644 index 0000000..b1968cf --- /dev/null +++ b/src/ternary_kernels_neon.cpp @@ -0,0 +1,107 @@ +#include "ternary_kernels.hpp" + +#if defined(__aarch64__) && defined(__ARM_FEATURE_DOTPROD) +#include + +#include +#if defined(__linux__) +#include +#ifndef HWCAP_ASIMDDP +#define HWCAP_ASIMDDP (1 << 20) +#endif +#endif + +namespace pk { +namespace { + +// One int32x4 lane per output row: a 16-row block is four vectors. sdot by +// lane multiplies the codes (0..2, as int8) of 4 rows by 4 elements with one +// broadcast group of 4 activations. Same math as ternary_matmul_rows_ref: +// per group an exact int32 sum, then acc = acc + scale * (float)(sum - group_sum) +// in group order (separate multiply and add), then acc * act_scale. +template +void tile_neon(const TernaryWeight& w, const uint8_t* act, size_t rb, int t0, float* y, int b0, int r0, + int r1) { + const int N = w.N, K = w.K, G = w.groups(); + const uint8x16_t m3 = vdupq_n_u8(3); + const int8_t* xr[C]; + const int32_t* gsr[C]; + for (int c = 0; c < C; ++c) { + const uint8_t* row = act + (size_t)(t0 + c) * rb; + xr[c] = reinterpret_cast(row); + gsr[c] = reinterpret_cast(row + K + 4); + } + float32x4_t facc[C][4]; + for (int c = 0; c < C; ++c) + for (int q = 0; q < 4; ++q) facc[c][q] = vdupq_n_f32(0.0f); + for (int g = 0; g < G; ++g) { + int32x4_t ia[C][4]; + for (int c = 0; c < C; ++c) + for (int q = 0; q < 4; ++q) ia[c][q] = vdupq_n_s32(0); + const uint8_t* pb = w.planes.data() + ((size_t)b0 * G + g) * 512; + for (int s = 0; s < 8; ++s) { + uint8x16_t v[4]; + for (int q = 0; q < 4; ++q) v[q] = vld1q_u8(pb + s * 64 + 16 * q); + int8x16_t xa[C]; + for (int c = 0; c < C; ++c) xa[c] = vld1q_s8(xr[c] + (size_t)g * 128 + s * 16); + for (int q = 0; q < 4; ++q) { + const int8x16_t w0 = vreinterpretq_s8_u8(vandq_u8(v[q], m3)); + const int8x16_t w1 = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(v[q], 2), m3)); + const int8x16_t w2 = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(v[q], 4), m3)); + const int8x16_t w3 = vreinterpretq_s8_u8(vshrq_n_u8(v[q], 6)); + for (int c = 0; c < C; ++c) { + ia[c][q] = vdotq_laneq_s32(ia[c][q], w0, xa[c], 0); + ia[c][q] = vdotq_laneq_s32(ia[c][q], w1, xa[c], 1); + ia[c][q] = vdotq_laneq_s32(ia[c][q], w2, xa[c], 2); + ia[c][q] = vdotq_laneq_s32(ia[c][q], w3, xa[c], 3); + } + } + } + const float* scp = &w.scales[((size_t)b0 * G + g) * kTernaryRowBlock]; + for (int c = 0; c < C; ++c) { + const int32x4_t gs = vdupq_n_s32(gsr[c][g]); + for (int q = 0; q < 4; ++q) { + const float32x4_t f = vcvtq_f32_s32(vsubq_s32(ia[c][q], gs)); + facc[c][q] = vaddq_f32(facc[c][q], vmulq_f32(vld1q_f32(scp + 4 * q), f)); + } + } + } + for (int c = 0; c < C; ++c) { + const float32x4_t sa = vdupq_n_f32(*reinterpret_cast(xr[c] + K)); + float out[16]; + for (int q = 0; q < 4; ++q) vst1q_f32(out + 4 * q, vmulq_f32(facc[c][q], sa)); + float* yr = y + (size_t)(t0 + c) * N; + const int lo = std::max(r0, b0 * kTernaryRowBlock); + const int hi = std::min(r1, (b0 + 1) * kTernaryRowBlock); + for (int n = lo; n < hi; ++n) yr[n] = out[n - b0 * kTernaryRowBlock]; + } +} + +void rows_neon(const TernaryWeight& w, const uint8_t* act, int T, float* y, int r0, int r1) { + if (r0 >= r1 || T <= 0) return; + constexpr int C = 2; + const size_t rb = ternary_act_row_bytes(w.K); + const int bl = (r1 - 1) / kTernaryRowBlock; + for (int b = r0 / kTernaryRowBlock; b <= bl; ++b) { + int t0 = 0; + for (; t0 + C <= T; t0 += C) tile_neon(w, act, rb, t0, y, b, r0, r1); + for (; t0 < T; ++t0) tile_neon<1>(w, act, rb, t0, y, b, r0, r1); + } +} + +} // namespace + +const TernaryKernel* ternary_kernel_neon() { + static const TernaryKernel k{"neon", rows_neon}; +#if defined(__linux__) + if (!(getauxval(AT_HWCAP) & HWCAP_ASIMDDP)) return nullptr; +#endif + return &k; +} + +} // namespace pk +#else +namespace pk { +const TernaryKernel* ternary_kernel_neon() { return nullptr; } +} // namespace pk +#endif diff --git a/src/ternary_kernels_x86.cpp b/src/ternary_kernels_x86.cpp new file mode 100644 index 0000000..3e7abe8 --- /dev/null +++ b/src/ternary_kernels_x86.cpp @@ -0,0 +1,309 @@ +#include "ternary_kernels.hpp" + +#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__)) +#include + +#include +#include + +namespace pk { +namespace { + +#define PK_TGT_AVX2 __attribute__((target("avx2"))) +#define PK_TGT_512 __attribute__((target("avx2,avx512f,avx512vl,avx512vnni"))) +#define PK_TGT_512F __attribute__((target("avx2,avx512f"))) + +// Both kernels compute the same math as ternary_matmul_rows_ref. One vector +// lane holds one output row of a 16-row block (see TernaryWeight), so the +// per-group integer sums of 16 rows come out in one vector without any +// horizontal reduction. For each output element the float update is exactly +// the reference sequence: acc = acc + scale * (float)(sum - group_sum), in group +// order, then acc * act_scale (separate multiply and add, no FMA). + +inline int32_t load_i32(const int8_t* p) { + int32_t v; + std::memcpy(&v, p, 4); + return v; +} + +// Row mask of block b for the output range [r0, r1). +inline unsigned block_mask(int b, int r0, int r1) { + const int lo = std::max(r0 - b * kTernaryRowBlock, 0); + const int hi = std::min(r1 - b * kTernaryRowBlock, kTernaryRowBlock); + return hi <= lo ? 0u : ((1u << hi) - 1u) & ~((1u << lo) - 1u); +} + +// --------------------------------------------------------------------------- +// AVX-512 VNNI: RB row blocks (16 rows each, one zmm) by C activation rows. +template +PK_TGT_512 void tile_512(const TernaryWeight& w, const uint8_t* act, size_t rb, int t0, float* y, + int b0, const unsigned* mk) { + const int N = w.N, K = w.K, G = w.groups(); + const __m512i m3 = _mm512_set1_epi8(3); + const int8_t* xr[C]; + const int32_t* gsr[C]; + for (int c = 0; c < C; ++c) { + const uint8_t* row = act + (size_t)(t0 + c) * rb; + xr[c] = reinterpret_cast(row); + gsr[c] = reinterpret_cast(row + K + 4); + } + __m512 facc[RB][C]; + for (int b = 0; b < RB; ++b) + for (int c = 0; c < C; ++c) facc[b][c] = _mm512_setzero_ps(); + for (int g = 0; g < G; ++g) { + __m512i ia[RB][C]; + for (int b = 0; b < RB; ++b) + for (int c = 0; c < C; ++c) ia[b][c] = _mm512_setzero_si512(); + const uint8_t* pb[RB]; + for (int b = 0; b < RB; ++b) pb[b] = w.planes.data() + ((size_t)(b0 + b) * G + g) * 512; + for (int s = 0; s < 8; ++s) { + __m512i wv[RB][4]; + for (int b = 0; b < RB; ++b) { + const __m512i v = _mm512_loadu_si512(pb[b] + s * 64); + wv[b][0] = _mm512_and_si512(v, m3); + wv[b][1] = _mm512_and_si512(_mm512_srli_epi32(v, 2), m3); + wv[b][2] = _mm512_and_si512(_mm512_srli_epi32(v, 4), m3); + wv[b][3] = _mm512_and_si512(_mm512_srli_epi32(v, 6), m3); + } + for (int c = 0; c < C; ++c) { + const int8_t* x = xr[c] + (size_t)g * 128 + s * 16; + for (int p = 0; p < 4; ++p) { + const __m512i a = _mm512_set1_epi32(load_i32(x + 4 * p)); + for (int b = 0; b < RB; ++b) ia[b][c] = _mm512_dpbusd_epi32(ia[b][c], wv[b][p], a); + } + } + } + for (int b = 0; b < RB; ++b) { + const __m512 sc = _mm512_loadu_ps(&w.scales[((size_t)(b0 + b) * G + g) * kTernaryRowBlock]); + for (int c = 0; c < C; ++c) { + const __m512 f = _mm512_cvtepi32_ps(_mm512_sub_epi32(ia[b][c], _mm512_set1_epi32(gsr[c][g]))); + facc[b][c] = _mm512_add_ps(facc[b][c], _mm512_mul_ps(sc, f)); + } + } + } + for (int c = 0; c < C; ++c) { + const __m512 sa = _mm512_set1_ps(*reinterpret_cast(xr[c] + K)); + for (int b = 0; b < RB; ++b) + _mm512_mask_storeu_ps(y + (size_t)(t0 + c) * N + (size_t)(b0 + b) * kTernaryRowBlock, + (__mmask16)mk[b], _mm512_mul_ps(facc[b][c], sa)); + } +} + +template +PK_TGT_512 void cols_512(const TernaryWeight& w, const uint8_t* act, size_t rb, int T, float* y, int b0, + const unsigned* mk) { + int t0 = 0; + for (; t0 + C <= T; t0 += C) tile_512(w, act, rb, t0, y, b0, mk); + for (; t0 < T; ++t0) tile_512(w, act, rb, t0, y, b0, mk); +} + +PK_TGT_512 void rows_512(const TernaryWeight& w, const uint8_t* act, int T, float* y, int r0, int r1) { + if (r0 >= r1 || T <= 0) return; + // 3 row blocks by 4 activation rows measured best on Zen 5 (12 int32 and + // 12 float accumulators); 2 by 6 and 4 by 4 are within a few percent. + constexpr int C = 4; + const size_t rb = ternary_act_row_bytes(w.K); + const int bl = (r1 - 1) / kTernaryRowBlock; + for (int b = r0 / kTernaryRowBlock; b <= bl;) { + unsigned mk[3]; + const int nb = std::min(3, bl - b + 1); + for (int i = 0; i < nb; ++i) mk[i] = block_mask(b + i, r0, r1); + if (nb == 3) cols_512<3, C>(w, act, rb, T, y, b, mk); + else if (nb == 2) cols_512<2, C>(w, act, rb, T, y, b, mk); + else cols_512<1, C>(w, act, rb, T, y, b, mk); + b += nb; + } +} + +// --------------------------------------------------------------------------- +// AVX2: one row block as two ymm halves (rows 0-7, rows 8-15) by C activation +// rows. maddubs gives int16 pairs of at most 2 * 2 * 127 = 508 in magnitude, +// and a group adds 32 of them per lane, at most 16256, so a whole group +// accumulates exactly in int16 before one widening madd. +template +PK_TGT_AVX2 void tile_avx2(const TernaryWeight& w, const uint8_t* act, size_t rb, int t0, float* y, int b0, + unsigned mk) { + const int N = w.N, K = w.K, G = w.groups(); + const __m256i m3 = _mm256_set1_epi8(3); + const __m256i ones = _mm256_set1_epi16(1); + const int8_t* xr[C]; + const int32_t* gsr[C]; + for (int c = 0; c < C; ++c) { + const uint8_t* row = act + (size_t)(t0 + c) * rb; + xr[c] = reinterpret_cast(row); + gsr[c] = reinterpret_cast(row + K + 4); + } + alignas(32) float facc[C][16] = {}; + for (int g = 0; g < G; ++g) { + __m256i ia[C][2]; + for (int c = 0; c < C; ++c) ia[c][0] = ia[c][1] = _mm256_setzero_si256(); + const uint8_t* pb = w.planes.data() + ((size_t)b0 * G + g) * 512; + for (int s = 0; s < 8; ++s) { + const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(pb + s * 64)); + const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(pb + s * 64 + 32)); + for (int p = 0; p < 4; ++p) { + const __m256i w0 = _mm256_and_si256(_mm256_srli_epi32(v0, 2 * p), m3); + const __m256i w1 = _mm256_and_si256(_mm256_srli_epi32(v1, 2 * p), m3); + for (int c = 0; c < C; ++c) { + const __m256i a = _mm256_set1_epi32(load_i32(xr[c] + (size_t)g * 128 + s * 16 + 4 * p)); + ia[c][0] = _mm256_add_epi16(ia[c][0], _mm256_maddubs_epi16(w0, a)); + ia[c][1] = _mm256_add_epi16(ia[c][1], _mm256_maddubs_epi16(w1, a)); + } + } + } + const float* scp = &w.scales[((size_t)b0 * G + g) * kTernaryRowBlock]; + const __m256 sc0 = _mm256_loadu_ps(scp), sc1 = _mm256_loadu_ps(scp + 8); + for (int c = 0; c < C; ++c) { + const __m256i gs = _mm256_set1_epi32(gsr[c][g]); + const __m256 f0 = _mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_madd_epi16(ia[c][0], ones), gs)); + const __m256 f1 = _mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_madd_epi16(ia[c][1], ones), gs)); + _mm256_store_ps(facc[c], _mm256_add_ps(_mm256_load_ps(facc[c]), _mm256_mul_ps(sc0, f0))); + _mm256_store_ps(facc[c] + 8, _mm256_add_ps(_mm256_load_ps(facc[c] + 8), _mm256_mul_ps(sc1, f1))); + } + } + for (int c = 0; c < C; ++c) { + const __m256 sa = _mm256_set1_ps(*reinterpret_cast(xr[c] + K)); + alignas(32) float out[16]; + _mm256_store_ps(out, _mm256_mul_ps(_mm256_load_ps(facc[c]), sa)); + _mm256_store_ps(out + 8, _mm256_mul_ps(_mm256_load_ps(facc[c] + 8), sa)); + float* yr = y + (size_t)(t0 + c) * N + (size_t)b0 * kTernaryRowBlock; + if (mk == 0xFFFFu) std::memcpy(yr, out, sizeof(out)); + else + for (int i = 0; i < kTernaryRowBlock; ++i) + if (mk >> i & 1u) yr[i] = out[i]; + } +} + +PK_TGT_AVX2 void rows_avx2(const TernaryWeight& w, const uint8_t* act, int T, float* y, int r0, int r1) { + if (r0 >= r1 || T <= 0) return; + constexpr int C = 2; + const size_t rb = ternary_act_row_bytes(w.K); + const int bl = (r1 - 1) / kTernaryRowBlock; + for (int b = r0 / kTernaryRowBlock; b <= bl; ++b) { + const unsigned mk = block_mask(b, r0, r1); + int t0 = 0; + for (; t0 + C <= T; t0 += C) tile_avx2(w, act, rb, t0, y, b, mk); + for (; t0 < T; ++t0) tile_avx2<1>(w, act, rb, t0, y, b, mk); + } +} + +// --------------------------------------------------------------------------- +// Activation quantization, byte-identical to ternary_quant_rows_ref. The max +// of |x| is exact in any order, and max_ps(v, m) returns m for a NaN v, as +// std::max(amax, NaN) keeps amax. The scale and its inverse are the same +// scalar float divisions. The product is the same float multiply; the clamp +// uses max_ps(f, -127) then min_ps(f, 127), which match std::max(-127.f, f) +// and std::min(127.f, f) including NaN; cvtps rounds to nearest even like +// lrintf in the default rounding mode. Group sums are exact integers. +PK_TGT_512F void quant_avx512(const float* x, int K, int t0, int t1, uint8_t* act) { + const int G = K / kTernaryGroup; + const size_t rb = ternary_act_row_bytes(K); + const __m512 lo = _mm512_set1_ps(-127.0f), hi = _mm512_set1_ps(127.0f); + for (int t = t0; t < t1; ++t) { + const float* xr = x + (size_t)t * K; + uint8_t* row = act + (size_t)t * rb; + __m512 m = _mm512_setzero_ps(); + for (int k = 0; k < K; k += 16) m = _mm512_max_ps(_mm512_abs_ps(_mm512_loadu_ps(xr + k)), m); + const float amax = _mm512_reduce_max_ps(m); + const float sc = amax > 0.0f ? amax / 127.0f : 0.0f; + const float inv = sc > 0.0f ? 1.0f / sc : 0.0f; + std::memcpy(row + K, &sc, 4); + const __m512 vinv = _mm512_set1_ps(inv); + for (int g = 0; g < G; ++g) { + __m512i sum = _mm512_setzero_si512(); + for (int j = 0; j < kTernaryGroup; j += 16) { + const int k = g * kTernaryGroup + j; + __m512 f = _mm512_mul_ps(_mm512_loadu_ps(xr + k), vinv); + f = _mm512_min_ps(_mm512_max_ps(f, lo), hi); + const __m512i v = _mm512_cvtps_epi32(f); + sum = _mm512_add_epi32(sum, v); + _mm_storeu_si128(reinterpret_cast<__m128i*>(row + k), _mm512_cvtepi32_epi8(v)); + } + const int32_t gs = _mm512_reduce_add_epi32(sum); + std::memcpy(row + K + 4 + 4 * (size_t)g, &gs, 4); + } + } +} + +PK_TGT_AVX2 void quant_avx2(const float* x, int K, int t0, int t1, uint8_t* act) { + const int G = K / kTernaryGroup; + const size_t rb = ternary_act_row_bytes(K); + const __m256 lo = _mm256_set1_ps(-127.0f), hi = _mm256_set1_ps(127.0f); + const __m256 absm = _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffff)); + const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + for (int t = t0; t < t1; ++t) { + const float* xr = x + (size_t)t * K; + uint8_t* row = act + (size_t)t * rb; + __m256 m = _mm256_setzero_ps(); + for (int k = 0; k < K; k += 8) m = _mm256_max_ps(_mm256_and_ps(_mm256_loadu_ps(xr + k), absm), m); + alignas(32) float mv[8]; + _mm256_store_ps(mv, m); + float amax = 0.0f; + for (float v : mv) amax = std::max(amax, v); + const float sc = amax > 0.0f ? amax / 127.0f : 0.0f; + const float inv = sc > 0.0f ? 1.0f / sc : 0.0f; + std::memcpy(row + K, &sc, 4); + const __m256 vinv = _mm256_set1_ps(inv); + for (int g = 0; g < G; ++g) { + __m256i sum = _mm256_setzero_si256(); + for (int j = 0; j < kTernaryGroup; j += 32) { + const int k = g * kTernaryGroup + j; + __m256i v[4]; + for (int u = 0; u < 4; ++u) { + __m256 f = _mm256_mul_ps(_mm256_loadu_ps(xr + k + 8 * u), vinv); + f = _mm256_min_ps(_mm256_max_ps(f, lo), hi); + v[u] = _mm256_cvtps_epi32(f); + sum = _mm256_add_epi32(sum, v[u]); + } + // values are within [-127, 127], so the saturating packs are exact + const __m256i p16a = _mm256_packs_epi32(v[0], v[1]); + const __m256i p16b = _mm256_packs_epi32(v[2], v[3]); + const __m256i p8 = _mm256_permutevar8x32_epi32(_mm256_packs_epi16(p16a, p16b), perm); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(row + k), p8); + } + __m128i s4 = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum, 1)); + s4 = _mm_add_epi32(s4, _mm_shuffle_epi32(s4, _MM_SHUFFLE(1, 0, 3, 2))); + s4 = _mm_add_epi32(s4, _mm_shuffle_epi32(s4, _MM_SHUFFLE(2, 3, 0, 1))); + const int32_t gs = _mm_cvtsi128_si32(s4); + std::memcpy(row + K + 4 + 4 * (size_t)g, &gs, 4); + } + } +} + +} // namespace + +const TernaryKernel* ternary_kernel_x86_vnni() { + static const TernaryKernel k{"vnni", rows_512}; + __builtin_cpu_init(); + const bool ok = __builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512vl") && + __builtin_cpu_supports("avx512vnni"); + return ok ? &k : nullptr; +} + +const TernaryKernel* ternary_kernel_x86_avx2() { + static const TernaryKernel k{"avx2", rows_avx2}; + __builtin_cpu_init(); + return __builtin_cpu_supports("avx2") ? &k : nullptr; +} + +const TernaryQuant* ternary_quant_x86_avx512() { + static const TernaryQuant q{"avx512", quant_avx512}; + __builtin_cpu_init(); + return __builtin_cpu_supports("avx512f") ? &q : nullptr; +} + +const TernaryQuant* ternary_quant_x86_avx2() { + static const TernaryQuant q{"avx2", quant_avx2}; + __builtin_cpu_init(); + return __builtin_cpu_supports("avx2") ? &q : nullptr; +} + +} // namespace pk +#else +namespace pk { +const TernaryKernel* ternary_kernel_x86_vnni() { return nullptr; } +const TernaryKernel* ternary_kernel_x86_avx2() { return nullptr; } +const TernaryQuant* ternary_quant_x86_avx512() { return nullptr; } +const TernaryQuant* ternary_quant_x86_avx2() { return nullptr; } +} // namespace pk +#endif diff --git a/src/transducer_batch.hpp b/src/transducer_batch.hpp index a207236..dcb32d1 100644 --- a/src/transducer_batch.hpp +++ b/src/transducer_batch.hpp @@ -7,8 +7,13 @@ namespace pk { // Batched greedy decode for N utterances. encs[n]: row-major [T[n], enc_hidden]. // durations empty -> RNNT (advance-by-1); non-empty -> TDT (advance-by-duration). -// Outputs per item: ids[n], and (if toks != nullptr) TokenInfo[n]. Produces -// output bit-identical to per-item rnnt_greedy / tdt_greedy. +// Outputs per item: ids[n], and (if toks != nullptr) TokenInfo[n]. Applies the +// same per-item decision rule as rnnt_greedy / tdt_greedy, but it is NOT +// guaranteed bit-identical to them: for N > 1 ggml picks a different matmul +// kernel than for N = 1, so prediction and joint outputs differ in float +// (logits by up to about 1e-4 measured on the packed Redux model). The emitted +// token sequences were identical on the test clips; a near-tie argmax could in +// principle flip. The tests compare with a tolerance. void transducer_greedy_batch( const PredictionNet& pred, const Joint& joint, const std::vector>& encs, diff --git a/src/vad_head.cpp b/src/vad_head.cpp new file mode 100644 index 0000000..2d9f125 --- /dev/null +++ b/src/vad_head.cpp @@ -0,0 +1,80 @@ +#include "vad_head.hpp" + +#include +#include +#include +#include + +#include "backend.hpp" +#include "ggml.h" +#include "model_loader.hpp" + +namespace pk { + +namespace { + +std::vector load_f32(const ModelLoader& ml, const std::string& name, size_t expect) { + const ggml_tensor* t = ml.tensor(name); + if (!t) throw std::runtime_error("VAD head: missing tensor " + name); + if (t->type != GGML_TYPE_F32 || (size_t)ggml_nelements(t) != expect) + throw std::runtime_error("VAD head: unexpected type or size for " + name); + // The tensor may live in device memory: read it through the backend. + std::vector out; + weight_to_host_f32(ml, name.c_str(), out); + return out; +} + +} // namespace + +VadHead::VadHead(const ModelLoader& ml) { + const VadCfg& c = ml.config().vad; + if (!c.present) throw std::runtime_error("model has no VAD head"); + w_.d_in = (int)c.d_in; + w_.hidden = (int)c.hidden; + w_.kernel = (int)c.kernel; + frame_sec_ = c.frame_sec; + const size_t D = c.d_in, H = c.hidden, K = c.kernel; + w_.proj_w = load_f32(ml, "vad_head.proj.weight", H * D); + w_.proj_b = load_f32(ml, "vad_head.proj.bias", H); + w_.ctx_w = load_f32(ml, "vad_head.ctx.weight", H * H * K); + w_.ctx_b = load_f32(ml, "vad_head.ctx.bias", H); + w_.out_w = load_f32(ml, "vad_head.out.weight", H); + w_.out_b = load_f32(ml, "vad_head.out.bias", 1)[0]; +} + +std::vector VadHead::run(const VadWeights& w, const VadVariant& v, const float* x, int T) { + const int D = w.d_in, H = w.hidden, K = w.kernel, pad = K / 2; + std::vector h1((size_t)T * H), p((size_t)T); + for (int t = 0; t < T; ++t) + for (int o = 0; o < H; ++o) { + float acc = w.proj_b[o]; + const float* wr = &w.proj_w[(size_t)o * D]; + const float* xr = x + (size_t)t * D; + for (int i = 0; i < D; ++i) acc += wr[i] * xr[i]; + h1[(size_t)t * H + o] = v.relu_after_proj ? std::max(0.0f, acc) : acc; + } + for (int t = 0; t < T; ++t) { + std::vector h2(H); + for (int o = 0; o < H; ++o) { + float acc = w.ctx_b[o]; + for (int kk = 0; kk < K; ++kk) { + const int tt = t + kk - pad; + if (tt < 0 || tt >= T) continue; + const float* hr = &h1[(size_t)tt * H]; + for (int i = 0; i < H; ++i) acc += w.ctx_w[((size_t)o * H + i) * K + kk] * hr[i]; + } + if (v.residual) acc += h1[(size_t)t * H + o]; + h2[o] = v.relu_after_ctx ? std::max(0.0f, acc) : acc; + } + float z = w.out_b; + for (int i = 0; i < H; ++i) z += w.out_w[i] * h2[i]; + p[t] = 1.0f / (1.0f + std::exp(-z)); + } + return p; +} + +std::vector VadHead::probabilities(const float* x, int T, const VadVariant* v) const { + return run(w_, v ? *v : VadVariant(), x, T); +} + +} // namespace pk diff --git a/src/vad_head.hpp b/src/vad_head.hpp new file mode 100644 index 0000000..cacfcaf --- /dev/null +++ b/src/vad_head.hpp @@ -0,0 +1,49 @@ +#pragma once +// Voice-activity head that moondream/parakeet-ultra and -redux carry on the +// encoder's subsampler output: proj (1x1 conv) -> ctx (conv, k=5) -> out (1x1 +// conv) -> sigmoid, one probability per 80 ms frame. +#include + +namespace pk { + +class ModelLoader; + +struct VadWeights { + int d_in = 0, hidden = 0, kernel = 0; + std::vector proj_w, proj_b; // [hidden][d_in], [hidden] + std::vector ctx_w, ctx_b; // [hidden][hidden][kernel], [hidden] + std::vector out_w; // [hidden] + float out_b = 0.0f; +}; + +// The three choices the checkpoint does not document. Fixed from evidence in +// docs/ternary.md (section "VAD head wiring"). +struct VadVariant { + bool relu_after_proj = true; + bool residual = false; + bool relu_after_ctx = true; + static VadVariant from_index(int i) { + VadVariant v; + v.relu_after_proj = (i & 1) != 0; + v.residual = (i & 2) != 0; + v.relu_after_ctx = (i & 4) != 0; + return v; + } +}; + +class VadHead { +public: + // Throws std::runtime_error if the GGUF has no VAD head or its tensors are + // missing or mis-shaped. + explicit VadHead(const ModelLoader& ml); + // x: row-major [T][d_in]. Returns T probabilities in [0, 1]. + std::vector probabilities(const float* x, int T, const VadVariant* v = nullptr) const; + float frame_sec() const { return frame_sec_; } + static std::vector run(const VadWeights& w, const VadVariant& v, const float* x, int T); + +private: + VadWeights w_; + float frame_sec_ = 0.08f; +}; + +} // namespace pk diff --git a/src/vad_segmenter.cpp b/src/vad_segmenter.cpp new file mode 100644 index 0000000..7f02848 --- /dev/null +++ b/src/vad_segmenter.cpp @@ -0,0 +1,66 @@ +#include "vad_segmenter.hpp" + +#include +#include +#include + +namespace pk { + +namespace { +// Above this many seconds a frame count could overflow int64. +constexpr double kMaxSec = 1e6; +} + +std::vector segment_by_vad(const std::vector& p, double total_sec, + const SegmenterOpts& o) { + std::vector out; + // Degenerate options: no usable frame grid, so do not cut at all. + if (!(o.frame_sec > 0.0) || !std::isfinite(o.frame_sec) || !std::isfinite(o.max_seg_sec) || + !(o.max_seg_sec > 2.0 * o.frame_sec) || !std::isfinite(total_sec) || + !std::isfinite(o.threshold) || !std::isfinite(o.min_seg_sec) || + !std::isfinite(o.min_pause_sec) || o.max_seg_sec > kMaxSec || o.min_seg_sec > kMaxSec || + o.min_pause_sec > kMaxSec) { + out.push_back({0.0, total_sec}); + return out; + } + if (total_sec <= o.max_seg_sec) { + out.push_back({0.0, total_sec}); + return out; + } + const double fs = o.frame_sec; + const int64_t n = (int64_t)p.size(); + const int64_t max_f = std::max(2, (int64_t)std::floor(o.max_seg_sec / fs + 1e-9)); + const int64_t min_f = std::min(max_f - 1, std::max(1, (int64_t)std::ceil(o.min_seg_sec / fs - 1e-9))); + const int64_t pause_f = std::max(1, (int64_t)std::ceil(o.min_pause_sec / fs - 1e-9)); + auto silent = [&](int64_t f) { return f >= n || p[(size_t)f] < o.threshold; }; + + // Middle of the longest silent run inside [lo, hi) that is at least pause_f + // frames long (ties go to the later run); -1 if there is none. + auto best_cut = [&](int64_t lo, int64_t hi) -> int64_t { + int64_t best = -1, best_len = 0; + int64_t f = lo; + while (f < hi) { + if (!silent(f)) { ++f; continue; } + int64_t e = f; + while (e < hi && silent(e)) ++e; + const int64_t len = e - f; + if (len >= pause_f && len >= best_len) { best_len = len; best = (f + e) / 2; } + f = e; + } + return best; + }; + + int64_t s = 0; + while (total_sec - (double)s * fs > o.max_seg_sec + 1e-9) { + const int64_t hi = s + max_f; + int64_t c = best_cut(s + (max_f * 2) / 3, hi); + if (c < 0) c = best_cut(s + min_f, hi); + if (c <= s) c = hi; // hard cut + out.push_back({(double)s * fs, (double)c * fs}); + s = c; + } + out.push_back({(double)s * fs, total_sec}); + return out; +} + +} // namespace pk diff --git a/src/vad_segmenter.hpp b/src/vad_segmenter.hpp new file mode 100644 index 0000000..36d346e --- /dev/null +++ b/src/vad_segmenter.hpp @@ -0,0 +1,31 @@ +#pragma once +#include + +namespace pk { + +struct VadSegment { + double start = 0.0; // seconds + double end = 0.0; +}; + +struct SegmenterOpts { + float threshold = 0.5f; // frame is speech when p >= threshold + double frame_sec = 0.08; + double max_seg_sec = 30.0; + double min_pause_sec = 0.32; // shortest silence that may be cut in + double min_seg_sec = 8.0; // do not cut earlier than this into a segment +}; + +// Cuts [0, total_sec] into contiguous segments of at most max_seg_sec, at +// pauses found in the per-frame speech probabilities p. Rule: cut in the middle +// of the longest run of frames below threshold (at least min_pause_sec long), +// searched first in the last third of the allowed window then in the whole +// window from min_seg_sec on; ties go to the later run; hard cut at max_seg_sec +// if no pause is found. Degenerate options (frame_sec not finite or <= 0, +// max_seg_sec not finite or <= 2 * frame_sec, threshold, min_seg_sec or +// min_pause_sec not finite, any of the three durations above 1e6 seconds) return the single segment +// {0, total_sec}. Every internal boundary is a whole number of frames. +std::vector segment_by_vad(const std::vector& p, double total_sec, + const SegmenterOpts& o); + +} // namespace pk diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b143508..7611803 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -12,6 +12,10 @@ 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) +pk_add_test(test_model_loader_ternary) +pk_add_test(test_ternary_load_negative) +pk_add_test(test_vad_head) +pk_add_test(test_vad_segmenter) pk_add_test(test_fft) pk_add_test(test_mel) pk_add_test(test_mel_gpu) @@ -40,6 +44,12 @@ pk_add_test(test_prediction_step_batch) pk_add_test(test_joint) pk_add_test(test_joint_step_batch) pk_add_test(test_prompt_kernel) +pk_add_test(test_ternary) +pk_add_test(test_ternary_model) +pk_add_test(test_ternary_long) +add_executable(bench_ternary bench_ternary.cpp) +target_link_libraries(bench_ternary PRIVATE parakeet) +target_include_directories(bench_ternary PRIVATE ${CMAKE_SOURCE_DIR}/src ${CMAKE_SOURCE_DIR}/third_party) pk_add_test(test_transducer_core) pk_add_test(test_tdt_greedy) pk_add_test(test_tdt_beam) @@ -48,6 +58,7 @@ pk_add_test(test_transducer_greedy_batch_rnnt) pk_add_test(test_timestamps_tokens) pk_add_test(test_timestamps) pk_add_test(test_transcribe_batch_ts) +pk_add_test(test_decoder_reuse) pk_add_test(test_tokenizer) pk_add_test(test_special_token_filter) pk_add_test(test_transcribe) @@ -59,6 +70,7 @@ pk_add_test(test_transcribe_ctc) pk_add_test(test_transcribe_rnnt) pk_add_test(test_transcribe_eou) pk_add_test(test_transcribe_nemotron) +pk_add_test(test_transcribe_vad) pk_add_test(test_streaming_decode) pk_add_test(test_streaming_eou_reset) pk_add_test(test_streaming_nemotron) @@ -147,6 +159,8 @@ set_tests_properties(test_model_loader test_mel test_mel_gpu test_subsampling te test_capi_timestamps test_capi_batch_json test_capi_ctc_logits test_combined_offline test_streaming_diarization test_diarization_accuracy test_diarization PROPERTIES LABELS "model") +set_tests_properties(test_ternary_model test_ternary_long test_decoder_reuse test_transcribe_vad test_vad_head test_model_loader_ternary + test_ternary_load_negative PROPERTIES LABELS "model") # These tests read fixtures/baselines via paths relative to the project root. set_tests_properties(test_mel test_mel_gpu test_subsampling test_subsampling_batch test_subsampling_batch_causal test_relpos_attention test_relpos_attention_batch test_conformer test_conformer_batch test_conv_eou test_encoder test_encoder_batch test_encoder_batch_local test_encoder_eou test_streaming_encoder @@ -162,7 +176,7 @@ set_tests_properties(test_mel test_mel_gpu test_subsampling test_subsampling_bat test_transcribe_ctc test_transcribe_rnnt test_transcribe_eou test_transcribe_nemotron test_streaming_decode test_streaming_eou_reset test_streaming_nemotron test_streaming_mel test_capi test_capi_batch test_capi_stream test_capi_stream_json test_capi_timestamps test_capi_batch_json test_capi_ctc_logits - test_combined_offline test_streaming_diarization test_diarization_accuracy + test_combined_offline test_streaming_diarization test_diarization_accuracy test_ternary_model test_ternary_long test_transcribe_vad test_decoder_reuse PROPERTIES WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}) # Python converter check (skips with exit 77 when the venv/model are absent). @@ -173,4 +187,7 @@ if(PARAKEET_PY) add_test(NAME check_baseline COMMAND ${PARAKEET_PY} ${CMAKE_SOURCE_DIR}/tests/python/check_baseline.py) set_tests_properties(check_baseline PROPERTIES SKIP_RETURN_CODE 77 LABELS "model") + + add_test(NAME check_hf_convert COMMAND ${PARAKEET_PY} ${CMAKE_SOURCE_DIR}/tests/python/check_hf_convert.py) + set_tests_properties(check_hf_convert PROPERTIES SKIP_RETURN_CODE 77) endif() diff --git a/tests/bench_ternary.cpp b/tests/bench_ternary.cpp new file mode 100644 index 0000000..b8d0c2e --- /dev/null +++ b/tests/bench_ternary.cpp @@ -0,0 +1,181 @@ +// Single-thread throughput of each ternary kernel, next to ggml's own mul_mat +// with Q8_0 and F16 weights of the same shape (1 thread, activations F32, so the +// ggml rows include their per-call activation quantization; the ternary rows do +// not include ternary_quant_rows, which is timed on its own). Usage: bench_ternary [N K T reps] +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ggml.h" +#include "ggml-cpu.h" +#include "ternary_kernels.hpp" + +using namespace pk; + +// Time ggml_mul_mat(W[K,N] of type wt, X[K,T] f32) on one thread; returns seconds per run. +static double time_ggml_mul_mat(ggml_type wt, int N, int K, int T, int reps, const std::vector& x, + std::mt19937& rng) { + ggml_init_params ip{}; + ip.mem_size = (size_t)N * K * 4 + (size_t)T * K * 4 + (size_t)T * N * 4 + (64u << 20); + ip.no_alloc = false; + ggml_context* ctx = ggml_init(ip); + ggml_tensor* w = ggml_new_tensor_2d(ctx, wt, K, N); + std::vector wf((size_t)N * K); + std::uniform_real_distribution d(-1.f, 1.f); + for (auto& v : wf) v = d(rng); + if (wt == GGML_TYPE_F32) std::memcpy(w->data, wf.data(), ggml_nbytes(w)); + else if (wt == GGML_TYPE_F16) ggml_fp32_to_fp16_row(wf.data(), (ggml_fp16_t*)w->data, (int64_t)N * K); + else ggml_quantize_chunk(wt, wf.data(), w->data, 0, N, K, nullptr); + ggml_tensor* xt = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, T); + std::memcpy(xt->data, x.data(), ggml_nbytes(xt)); + ggml_tensor* y = ggml_mul_mat(ctx, w, xt); + ggml_cgraph* g = ggml_new_graph(ctx); + ggml_build_forward_expand(g, y); + ggml_cplan plan = ggml_graph_plan(g, 1, nullptr); + std::vector work(plan.work_size); + plan.work_data = work.data(); + ggml_graph_compute(g, &plan); // warm up + const auto t0 = std::chrono::steady_clock::now(); + for (int r = 0; r < reps; ++r) ggml_graph_compute(g, &plan); + const double sec = std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / reps; + ggml_free(ctx); + return sec; +} + +// repack mode: bench_ternary repack [weights=264]. Times the load-time repack of +// a Redux-sized weight set (alternating 4096x1024 and 1024x4096) with the old +// scalar reference, the new serial repack, and the new repack across a thread pool. +static void repack_reference_bench(const uint8_t* q, const uint16_t* s, int N, int K, TernaryWeight& out) { + const int G = K / kTernaryGroup, nb = (K + 4) / 5; + out.N = N; + out.K = K; + const int B = out.row_blocks(); + out.planes.assign((size_t)B * G * 512, 0); + out.scales.assign((size_t)B * G * kTernaryRowBlock, 0.0f); + std::vector row(K); + for (int n = 0; n < N; ++n) { + const int b = n / kTernaryRowBlock, i = n % kTernaryRowBlock; + for (int g = 0; g < G; ++g) + out.scales[((size_t)b * G + g) * kTernaryRowBlock + i] = ggml_fp16_to_fp32(s[(size_t)n * G + g]); + const uint8_t* qr = q + (size_t)n * nb; + for (int bb = 0; bb < nb; ++bb) { + int v = qr[bb]; + for (int d = 0; d < 5; ++d) { + const int k = bb * 5 + d; + if (k < K) row[k] = (uint8_t)(v % 3); + v /= 3; + } + } + for (int k = 0; k < K; ++k) { + const int g = k / kTernaryGroup, r = k % kTernaryGroup; + const int st = r / 16, p = (r % 16) / 4, j = r % 4; + out.planes[((size_t)b * G + g) * 512 + (size_t)st * 64 + 4 * i + j] |= (uint8_t)(row[k] << (2 * p)); + } + } +} + +static int repack_bench(int total) { + std::mt19937 rng(3); + const int shapes[2][2] = {{4096, 1024}, {1024, 4096}}; + std::vector q[2]; + std::vector s[2]; + for (int k = 0; k < 2; ++k) { + const int N = shapes[k][0], K = shapes[k][1]; + q[k].resize((size_t)N * ((K + 4) / 5)); + for (auto& b : q[k]) b = (uint8_t)(rng() % 243); + s[k].assign((size_t)N * (K / kTernaryGroup), ggml_fp32_to_fp16(0.05f)); + } + using clk = std::chrono::steady_clock; + auto ms = [](clk::time_point a) { return std::chrono::duration(clk::now() - a).count(); }; + std::vector out(total); + for (int k = 0; k < 2; ++k) { + TernaryWeight w; + auto t0 = clk::now(); + repack_reference_bench(q[k].data(), s[k].data(), shapes[k][0], shapes[k][1], w); + const double old_ms = ms(t0); + TernaryWeight w2; + t0 = clk::now(); + ternary_repack(q[k].data(), s[k].data(), shapes[k][0], shapes[k][1], w2); + std::printf("per weight %dx%d: old %.2f ms, new %.2f ms, identical=%d\n", shapes[k][0], shapes[k][1], old_ms, + ms(t0), (int)(w.planes == w2.planes)); + } + auto run = [&](bool old, int nt) { + std::atomic next{0}; + auto fn = [&] { + for (int i; (i = next.fetch_add(1)) < total;) { + const int k = i & 1; + if (old) repack_reference_bench(q[k].data(), s[k].data(), shapes[k][0], shapes[k][1], out[i]); + else ternary_repack(q[k].data(), s[k].data(), shapes[k][0], shapes[k][1], out[i]); + } + }; + const auto t0 = clk::now(); + std::vector th; + for (int t = 1; t < nt; ++t) th.emplace_back(fn); + fn(); + for (auto& t : th) t.join(); + return ms(t0); + }; + auto report = [&](const char* name, bool old, int nt) { + std::vector v; + for (int r = 0; r < 5; ++r) v.push_back(run(old, nt)); + std::sort(v.begin(), v.end()); + std::printf("total %d weights, %-22s threads=%d: min %.1f ms, median %.1f ms\n", total, name, nt, v[0], v[2]); + }; + report("old reference", true, 1); + report("new serial", false, 1); + report("new pool", false, 4); + report("new pool", false, 8); + return 0; +} + +int main(int argc, char** argv) { + if (argc > 1 && !std::strcmp(argv[1], "repack")) return repack_bench(argc > 2 ? std::atoi(argv[2]) : 264); + const int N = argc > 1 ? std::atoi(argv[1]) : 4096; + const int K = argc > 2 ? std::atoi(argv[2]) : 1024; + const int T = argc > 3 ? std::atoi(argv[3]) : 200; + const int reps = argc > 4 ? std::atoi(argv[4]) : 20; + std::mt19937 rng(1); + const int nb = (K + 4) / 5, G = K / kTernaryGroup; + std::vector q((size_t)N * nb); + for (auto& b : q) b = (uint8_t)(rng() % 243); + std::vector s((size_t)N * G, ggml_fp32_to_fp16(0.05f)); + TernaryWeight w; + ternary_repack(q.data(), s.data(), N, K, w); + std::vector x((size_t)T * K); + std::uniform_real_distribution d(-1.f, 1.f); + for (auto& v : x) v = d(rng); + std::vector act(ternary_act_row_bytes(K) * T); + ternary_quant_rows(x.data(), K, 0, T, act.data()); + std::vector y((size_t)T * N); + { + // activation quantization (op 1 of ternary_linear), per call + ternary_quant_rows(x.data(), K, 0, T, act.data()); + const auto t0 = std::chrono::steady_clock::now(); + for (int r = 0; r < reps; ++r) ternary_quant_rows(x.data(), K, 0, T, act.data()); + const double sec = std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / reps; + std::printf("quant N=- K=%d T=%d %8.3f ms (ternary_quant_rows, not in the kernel rows)\n", K, T, + sec * 1e3); + } + for (const TernaryKernel* k : ternary_all_kernels()) { + // the scalar reference is 50x slower; a few runs are enough for it + const int kr = std::strcmp(k->name, "scalar") ? reps : std::max(1, reps / 10); + k->fn(w, act.data(), T, y.data(), 0, N); // warm up + const auto t0 = std::chrono::steady_clock::now(); + for (int r = 0; r < kr; ++r) k->fn(w, act.data(), T, y.data(), 0, N); + const double sec = std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / kr; + std::printf("%-8s N=%d K=%d T=%d %8.3f ms %7.2f GMAC/s\n", k->name, N, K, T, sec * 1e3, + (double)N * K * T / sec / 1e9); + } + for (ggml_type wt : {GGML_TYPE_Q8_0, GGML_TYPE_F16}) { + const double sec = time_ggml_mul_mat(wt, N, K, T, reps, x, rng); + std::printf("ggml_%-3s N=%d K=%d T=%d %8.3f ms %7.2f GMAC/s\n", ggml_type_name(wt), N, K, T, sec * 1e3, + (double)N * K * T / sec / 1e9); + } + return 0; +} diff --git a/tests/python/check_hf_convert.py b/tests/python/check_hf_convert.py new file mode 100644 index 0000000..1e3ae7f --- /dev/null +++ b/tests/python/check_hf_convert.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +"""Model-independent checks for scripts/convert_hf_parakeet_to_gguf.py.""" +import importlib.util +import pathlib +import sys + +try: + import numpy as np + spec = importlib.util.spec_from_file_location( + "hfconv", + pathlib.Path(__file__).resolve().parents[2] / "scripts" / "convert_hf_parakeet_to_gguf.py") + hfconv = importlib.util.module_from_spec(spec) + spec.loader.exec_module(hfconv) +except SystemExit: + sys.exit(77) # gguf missing +except ImportError: + sys.exit(77) + + +def pack(codes): + """Inverse of unpack_ternary: codes [out, in] in {0,1,2} -> uint8 [out, ceil(in/5)].""" + out, n = codes.shape + nb = -(-n // 5) + pad = np.zeros((out, nb * 5), dtype=np.int64) + pad[:, :n] = codes + return (pad.reshape(out, nb, 5) * (3 ** np.arange(5))).sum(-1).astype(np.uint8) + + +def test_unpack_matches_loop(n_in, group=128, out=7, seed=0): + rng = np.random.default_rng(seed) + codes = rng.integers(0, 3, size=(out, n_in)) + ng = -(-n_in // group) + scales = rng.random((out, ng)).astype(np.float16) + got = hfconv.unpack_ternary(pack(codes), scales, n_in, group) + want = np.zeros((out, n_in), dtype=np.float32) + for r in range(out): + for c in range(n_in): + want[r, c] = np.float32(scales[r, c // group]) * (codes[r, c] - 1) + assert got.shape == (out, n_in) and got.dtype == np.float32 + assert np.array_equal(got, want), f"unpack mismatch for in={n_in}" + + +def test_unpack_rejects_bad_byte(): + q = np.array([[243]], dtype=np.uint8) # 3**5, not a valid 5-trit byte + try: + hfconv.unpack_ternary(q, np.ones((1, 1), dtype=np.float16), 5, 128) + except ValueError: + return + raise AssertionError("byte 243 must be rejected") + + +RENAMES = { + "encoder.subsampling.layers.3.weight": "encoder.pre_encode.conv.3.weight", + "encoder.subsampling.linear.bias": "encoder.pre_encode.out.bias", + "encoder.layers.5.conv.norm.running_var": "encoder.layers.5.conv.batch_norm.running_var", + "encoder.layers.5.self_attn.q_proj.weight": "encoder.layers.5.self_attn.linear_q.weight", + "encoder.layers.5.self_attn.k_proj.weight": "encoder.layers.5.self_attn.linear_k.weight", + "encoder.layers.5.self_attn.v_proj.weight": "encoder.layers.5.self_attn.linear_v.weight", + "encoder.layers.5.self_attn.o_proj.weight": "encoder.layers.5.self_attn.linear_out.weight", + "encoder.layers.5.self_attn.relative_k_proj.weight": "encoder.layers.5.self_attn.linear_pos.weight", + "encoder.layers.5.self_attn.bias_u": "encoder.layers.5.self_attn.pos_bias_u", + "encoder.layers.23.self_attn.bias_v": "encoder.layers.23.self_attn.pos_bias_v", + "decoder.embedding.weight": "decoder.prediction.embed.weight", + "decoder.lstm.weight_ih_l1": "decoder.prediction.dec_rnn.lstm.weight_ih_l1", + "encoder_projector.weight": "joint.enc.weight", + "decoder.decoder_projector.bias": "joint.pred.bias", + "joint.head.weight": "joint.joint_net.2.weight", + "encoder.layers.0.feed_forward1.linear1.weight": "encoder.layers.0.feed_forward1.linear1.weight", + "encoder.layers.0.norm_out.bias": "encoder.layers.0.norm_out.bias", + "encoder.layers.0.conv.pointwise_conv1.weight": "encoder.layers.0.conv.pointwise_conv1.weight", +} + + +def test_renames(): + for hf, nemo in RENAMES.items(): + assert hfconv.hf_to_nemo(hf) == nemo, (hf, hfconv.hf_to_nemo(hf), nemo) + try: + hfconv.hf_to_nemo("something.unknown") + except KeyError: + return + raise AssertionError("unknown tensor must raise KeyError") + + +def test_keep_matches_dequant(): + """Needs a real Redux download; skipped unless PK_HF_REDUX_DIR and PK_HF_TEMPLATE are set.""" + import os + import subprocess + import tempfile + d, t = os.environ.get("PK_HF_REDUX_DIR"), os.environ.get("PK_HF_TEMPLATE") + if not d or not t: + print("check_hf_convert: keep-vs-dequant skipped (set PK_HF_REDUX_DIR, PK_HF_TEMPLATE)") + return + import gguf + script = str(pathlib.Path(__file__).resolve().parents[2] / "scripts" / "convert_hf_parakeet_to_gguf.py") + with tempfile.TemporaryDirectory() as tmp: + keep, deq = tmp + "/keep.gguf", tmp + "/deq.gguf" + common = [sys.executable, script, "--hf", d, "--template", t, "--dtype", "f32"] + subprocess.run(common + ["--output", keep, "--ternary", "keep"], check=True) + subprocess.run(common + ["--output", deq, "--ternary", "dequant"], check=True) + ra, rb = gguf.GGUFReader(keep), gguf.GGUFReader(deq) + ta = {x.name: x for x in ra.tensors} + tb = {x.name: x for x in rb.tensors} + qn = sorted(n for n in ta if n.endswith(".qweight")) + assert len(qn) == 264, len(qn) + assert bool(ra.fields["parakeet.ternary.present"].contents()) + assert int(ra.fields["parakeet.ternary.group_size"].contents()) == 128 + assert bool(ra.fields["parakeet.vad.present"].contents()) + for n in ("vad_head.proj.weight", "vad_head.ctx.weight", "vad_head.out.weight", + "vad_head.proj.bias", "vad_head.ctx.bias", "vad_head.out.bias"): + assert n in ta and n in tb, n + for n in qn[:11] + qn[-11:]: + base = n[: -len(".qweight")] + assert ta[n].tensor_type.name == "I8", ta[n].tensor_type + assert (base + ".weight") not in ta + qw = np.array(ta[n].data).view(np.uint8) + sc = np.array(ta[base + ".scales"].data) + want = np.array(tb[base + ".weight"].data) + want = want.reshape(want.shape[0], -1) + got = hfconv.unpack_ternary(qw, sc, want.shape[1], 128) + assert np.array_equal(got, want), base + print("check_hf_convert: keep-vs-dequant OK") + + +if __name__ == "__main__": + for n in (1024, 4096, 130, 5, 128): + test_unpack_matches_loop(n) + test_unpack_rejects_bad_byte() + test_renames() + print("check_hf_convert: OK") + test_keep_matches_dequant() diff --git a/tests/test_capi_batch.cpp b/tests/test_capi_batch.cpp index 228447c..c4cbd16 100644 --- a/tests/test_capi_batch.cpp +++ b/tests/test_capi_batch.cpp @@ -9,7 +9,8 @@ // // Loads the real-speech fixture, transcribes it single-clip, then runs it as a // 2-clip batch of the same audio. Asserts both batch results equal the single -// result, proving the batch path is byte-identical to the single path. This is +// result, proving the batch path gives the same transcript as the single path +// (text equality; batched floats are close, not bit-identical). This is // self-consistency (our own code on both sides), so it only needs // PARAKEET_TEST_GGUF. int main() { diff --git a/tests/test_capi_batch_json.cpp b/tests/test_capi_batch_json.cpp index 038603d..33f9bb2 100644 --- a/tests/test_capi_batch_json.cpp +++ b/tests/test_capi_batch_json.cpp @@ -67,7 +67,7 @@ int main() { // Prompt (multilingual / nemotron) model: exercise the batched target_lang // variant. This fixture is a CAUSAL streaming prompt model // (causal_downsampling=True). Batched causal subsampling is now supported - // (byte-identical to per-item), so a valid-language 2-clip batch runs through + // (same transcripts as per-item), so a valid-language 2-clip batch runs through // the batched encoder and returns a JSON array of length 2. We also assert // the catchable error path: an unknown locale is rejected by // resolve_prompt_index (which runs before the encoder) -> NULL + non-empty diff --git a/tests/test_decoder_reuse.cpp b/tests/test_decoder_reuse.cpp new file mode 100644 index 0000000..d2112b6 --- /dev/null +++ b/tests/test_decoder_reuse.cpp @@ -0,0 +1,115 @@ +// The transducer decoder objects (PredictionNet, Joint) are built once per Model +// and reused across utterances. Reusing them must not change any result: +// - two consecutive transcribes of the same clip agree exactly (text, token +// ids, frames, spans and confidences), +// - a fresh Model (fresh decoder objects) agrees exactly with the reused one, +// - transcribe_pcm_batch agrees on text and token ids/frames (batched matmuls +// are only close in float, so confidences are compared with a tolerance), +// - two threads decoding on two different Model instances both work. +// Models: PARAKEET_TEST_GGUF_REDUX_KEEP and/or PARAKEET_TEST_GGUF_ULTRA. +// Skips (77) when neither is set. +#include "model.hpp" +#include "audio_io.hpp" +#include +#include +#include +#include +#include +#include + +// Batched encode (padded, tiled matmuls) is only close to the per-clip encode in +// float, so confidences move by about 0.01 on the packed ternary model; token +// ids, frames and spans are compared exactly. +static const float kBatchConfTol = 0.05f; + +static bool same_tokens(const std::vector& a, + const std::vector& b, float conf_tol) { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) { + if (a[i].id != b[i].id || a[i].frame != b[i].frame || a[i].span != b[i].span) + return false; + if (conf_tol == 0.0f ? a[i].conf != b[i].conf + : std::fabs(a[i].conf - b[i].conf) > conf_tol) + return false; + } + return true; +} + +static bool check_model(const char* path) { + std::fprintf(stderr, "model %s\n", path); + auto model = pk::Model::load(path); + auto model2 = pk::Model::load(path); + if (!model || !model2) { std::fprintf(stderr, "load failed\n"); return false; } + pk::Audio a; + if (!pk::load_audio_16k_mono("tests/fixtures/speech.wav", a) || a.samples.empty()) { + std::fprintf(stderr, "wav load failed\n"); return false; + } + std::vector half(a.samples.begin(), a.samples.begin() + (a.samples.size() * 3) / 4); + + // A fresh Model gives the reference (its decoder objects are used once). + pk::Transcription ref = model2->transcribe_with_timestamps(a.samples, 16000); + pk::Transcription ref2 = model2->transcribe_with_timestamps(half, 16000); + bool ok = !ref.tokens.empty(); + + // Reused decoder objects: same clip twice, interleaved with another clip. + pk::Transcription r1 = model->transcribe_with_timestamps(a.samples, 16000); + pk::Transcription r2 = model->transcribe_with_timestamps(half, 16000); + pk::Transcription r3 = model->transcribe_with_timestamps(a.samples, 16000); + ok = ok && r1.text == ref.text && same_tokens(r1.tokens, ref.tokens, 0.0f); + ok = ok && r3.text == ref.text && same_tokens(r3.tokens, ref.tokens, 0.0f); + ok = ok && r2.text == ref2.text && same_tokens(r2.tokens, ref2.tokens, 0.0f); + std::fprintf(stderr, " repeat/fresh: %s (%zu tokens)\n", ok ? "OK" : "FAIL", ref.tokens.size()); + + // The plain-text entry point shares the same objects. + ok = ok && model->transcribe_pcm(a.samples, 16000) == ref.text; + + // Batched path with the same reused objects. + auto b = model->transcribe_pcm_batch_with_timestamps({a.samples, half}, 16000); + if (b.size() != 2) { std::fprintf(stderr, " batch size %zu\n", b.size()); return false; } + bool bok = b[0].text == ref.text && b[1].text == ref2.text + && same_tokens(b[0].tokens, ref.tokens, kBatchConfTol) + && same_tokens(b[1].tokens, ref2.tokens, kBatchConfTol); + auto bt = model->transcribe_pcm_batch({a.samples, half}, 16000); + bok = bok && bt.size() == 2 && bt[0] == ref.text && bt[1] == ref2.text; + if (!bok) { + for (int i = 0; i < 2; ++i) { + const pk::Transcription& r = i ? ref2 : ref; + std::fprintf(stderr, " item %d: text %s, tokens %zu vs %zu\n batch: %s\n ref : %s\n", + i, b[i].text == r.text ? "same" : "DIFF", b[i].tokens.size(), + r.tokens.size(), b[i].text.c_str(), r.text.c_str()); + for (size_t k = 0; k < b[i].tokens.size() && k < r.tokens.size(); ++k) { + const pk::TokenInfo &x = b[i].tokens[k], &y = r.tokens[k]; + if (x.id != y.id || x.frame != y.frame || x.span != y.span || + std::fabs(x.conf - y.conf) > kBatchConfTol) + std::fprintf(stderr, " tok %zu: id %d/%d frame %d/%d span %d/%d conf %g/%g\n", + k, x.id, y.id, x.frame, y.frame, x.span, y.span, x.conf, y.conf); + } + } + } + std::fprintf(stderr, " batch: %s\n", bok ? "OK" : "FAIL"); + ok = ok && bok; + + // Two threads, two different Model instances. + pk::Transcription ta, tb; + std::thread t1([&] { ta = model->transcribe_with_timestamps(a.samples, 16000); }); + std::thread t2([&] { tb = model2->transcribe_with_timestamps(a.samples, 16000); }); + t1.join(); t2.join(); + bool tok = ta.text == ref.text && tb.text == ref.text + && same_tokens(ta.tokens, ref.tokens, 0.0f) && same_tokens(tb.tokens, ref.tokens, 0.0f); + std::fprintf(stderr, " two threads, two models: %s\n", tok ? "OK" : "FAIL"); + return ok && tok; +} + +int main() { + const char* envs[] = {"PARAKEET_TEST_GGUF_REDUX_KEEP", "PARAKEET_TEST_GGUF_ULTRA"}; + int ran = 0; + bool ok = true; + for (const char* e : envs) { + const char* p = std::getenv(e); + if (!p) continue; + ++ran; + ok = check_model(p) && ok; + } + if (!ran) { std::fprintf(stderr, "env not set; skip\n"); return 77; } + return ok ? 0 : 1; +} diff --git a/tests/test_model_loader_ternary.cpp b/tests/test_model_loader_ternary.cpp new file mode 100644 index 0000000..f8258c4 --- /dev/null +++ b/tests/test_model_loader_ternary.cpp @@ -0,0 +1,52 @@ +// Loader reads parakeet.ternary.* / parakeet.vad.* and leaves other models alone. +// Skips (77) unless the env vars point at converted GGUFs. +#include "model_loader.hpp" + +#include +#include +#include + +#include "ggml.h" + +using namespace pk; + +static int failures = 0; +#define CHECK(c) do { if (!(c)) { std::fprintf(stderr, "FAIL: %s (line %d)\n", #c, __LINE__); ++failures; } } while (0) + +int main() { + const char* keep = std::getenv("PARAKEET_TEST_GGUF_REDUX_KEEP"); + const char* plain = std::getenv("PARAKEET_TEST_GGUF"); + if (!keep && !plain) { + std::puts("skip: PARAKEET_TEST_GGUF_REDUX_KEEP and PARAKEET_TEST_GGUF unset"); + return 77; + } + if (keep) { + ModelLoader ml; + CHECK(ml.load(keep)); + const ParakeetConfig& c = ml.config(); + CHECK(c.ternary.present); + CHECK(c.ternary.group_size == 128); + CHECK(c.vad.present); + CHECK(c.vad.d_in == 1024 && c.vad.hidden == 128 && c.vad.kernel == 5); + CHECK(std::fabs(c.vad.frame_sec - 0.08f) < 1e-6f); + const ggml_tensor* q = ml.tensor("encoder.layers.0.self_attn.linear_q.qweight"); + CHECK(q && q->type == GGML_TYPE_I8 && q->ne[0] == 205 && q->ne[1] == 1024); + const ggml_tensor* s = ml.tensor("encoder.layers.0.self_attn.linear_q.scales"); + CHECK(s && s->type == GGML_TYPE_F16 && s->ne[0] == 8 && s->ne[1] == 1024); + CHECK(ml.tensor("encoder.layers.0.self_attn.linear_q.weight") == nullptr); + const ggml_tensor* p = ml.tensor("encoder.layers.0.feed_forward1.linear2.qweight"); + CHECK(p && p->ne[0] == 820 && p->ne[1] == 1024); // K = 4096, ceil(4096/5) = 820 + CHECK(ml.tensor("vad_head.proj.weight") != nullptr); + CHECK(ml.tensor("vad_head.ctx.weight") != nullptr); + CHECK(ml.tensor("vad_head.out.bias") != nullptr); + } + if (plain) { + ModelLoader ml; + CHECK(ml.load(plain)); + CHECK(!ml.config().ternary.present); + CHECK(!ml.config().vad.present); + } + if (failures) return 1; + std::puts("test_model_loader_ternary: OK"); + return 0; +} diff --git a/tests/test_subsampling_batch_causal.cpp b/tests/test_subsampling_batch_causal.cpp index d9b79fe..cc34e0e 100644 --- a/tests/test_subsampling_batch_causal.cpp +++ b/tests/test_subsampling_batch_causal.cpp @@ -8,7 +8,7 @@ // Batched CAUSAL subsampling parity for the multilingual streaming nemotron // model (causal_downsampling=True). Greedy decode is deterministic and the // per-item causal path is NeMo-validated at WER 0, so the gold check is -// byte-identical equivalence: a clip transcribed inside a B>1 batch MUST equal +// transcript equality: a clip transcribed inside a B>1 batch MUST equal // the SAME clip transcribed standalone. Mixed-length batches exercise the // per-item trailing-pad masking that interacts with the causal right pad of 1. // @@ -60,7 +60,8 @@ int main() { } // (b) MIXED-LENGTH batch (the real test): different lengths exercise the - // per-item causal masking; each item must be byte-identical to standalone. + // per-item causal masking; each item must give the same transcript as standalone + // (text equality; batched floats are close, not bit-identical). { auto out = model->transcribe_pcm_batch({speech.samples, clip.samples}, 16000, pk::Decoder::kDefault, lang); @@ -90,8 +91,8 @@ int main() { } // (b'') MIXED-LENGTH with a non-empty shorter item, both orderings. The half - // slice is the padded/masked item in one ordering; its tokens must be - // byte-identical to its standalone transcript. + // slice is the padded/masked item in one ordering; its tokens must + // give the same transcript as its standalone run. { auto out1 = model->transcribe_pcm_batch({speech.samples, half}, 16000, pk::Decoder::kDefault, lang); diff --git a/tests/test_ternary.cpp b/tests/test_ternary.cpp new file mode 100644 index 0000000..8808ee3 --- /dev/null +++ b/tests/test_ternary.cpp @@ -0,0 +1,311 @@ +// Unit test for the ternary core: repack, activation quantization, scalar +// reference. Model-independent. +#include "ternary.hpp" +#include "ternary_kernels.hpp" + +#include +#include +#include +#include +#include +#include + +#include "ggml.h" + +using namespace pk; + +static int failures = 0; +#define CHECK(c) do { if (!(c)) { std::fprintf(stderr, "FAIL: %s (line %d)\n", #c, __LINE__); ++failures; } } while (0) + +// Upstream packing: 5 trits per byte, base 3, least significant digit first. +static std::vector pack_trits(const std::vector& codes, int N, int K) { + const int nb = (K + 4) / 5; + std::vector q((size_t)N * nb, 0); + for (int n = 0; n < N; ++n) + for (int b = 0; b < nb; ++b) { + int v = 0, mul = 1; + for (int d = 0; d < 5; ++d) { + const int i = b * 5 + d; + const int c = i < K ? codes[(size_t)n * K + i] : 0; + v += c * mul; + mul *= 3; + } + q[(size_t)n * nb + b] = (uint8_t)v; + } + return q; +} + +// wide: codes mostly 2 and positive activations, so group sums reach 14 bits +// and scale * sum no longer fits a float mantissa exactly. Only then does a +// fused multiply-add round differently from the reference's multiply and add. +struct Case { int N, K, T; bool wide = false; }; + +static void run_case(const Case& cs, unsigned seed) { + std::mt19937 rng(seed); + std::uniform_int_distribution code(0, 2); + std::uniform_real_distribution sc(0.01f, 0.5f), xr(-3.0f, 3.0f); + const int N = cs.N, K = cs.K, T = cs.T, G = K / kTernaryGroup; + + std::vector codes((size_t)N * K); + std::uniform_int_distribution pct(0, 99); + for (auto& c : codes) c = cs.wide ? (pct(rng) < 90 ? 2 : code(rng)) : code(rng); + std::vector scales_f16((size_t)N * G); + std::vector scales_f((size_t)N * G); + for (size_t i = 0; i < scales_f16.size(); ++i) { + scales_f16[i] = ggml_fp32_to_fp16(sc(rng)); + scales_f[i] = ggml_fp16_to_fp32(scales_f16[i]); + } + const std::vector q = pack_trits(codes, N, K); + + // dequant equals the definition + std::vector W; + ternary_dequant(q.data(), scales_f16.data(), N, K, W); + CHECK(W.size() == (size_t)N * K); + for (int n = 0; n < N; ++n) + for (int k = 0; k < K; ++k) + CHECK(W[(size_t)n * K + k] == scales_f[(size_t)n * G + k / kTernaryGroup] * (float)(codes[(size_t)n * K + k] - 1)); + + TernaryWeight w; + ternary_repack(q.data(), scales_f16.data(), N, K, w); + CHECK(w.N == N && w.K == K && w.groups() == G); + CHECK(w.planes.size() == (size_t)w.row_blocks() * G * 512); + CHECK(w.scales.size() == (size_t)w.row_blocks() * G * kTernaryRowBlock); + // the repacked layout holds exactly the input codes and scales + for (int n = 0; n < N; ++n) { + for (int k = 0; k < K; ++k) CHECK(w.code(n, k) == codes[(size_t)n * K + k]); + for (int g = 0; g < G; ++g) CHECK(w.scale(n, g) == scales_f[(size_t)n * G + g]); + } + + std::vector x((size_t)T * K); + std::uniform_real_distribution xw(1.0f, 3.0f); + for (auto& v : x) v = cs.wide ? xw(rng) : xr(rng); + std::vector act(ternary_act_row_bytes(K) * T); + ternary_quant_rows(x.data(), K, 0, T, act.data()); + + // ref matmul over the full row range, then over two halves: identical + std::vector y((size_t)T * N), y2((size_t)T * N, -1.0f); + ternary_matmul_rows_ref(w, act.data(), T, y.data(), 0, N); + const int mid = N / 2; + ternary_matmul_rows_ref(w, act.data(), T, y2.data(), 0, mid); + ternary_matmul_rows_ref(w, act.data(), T, y2.data(), mid, N); + for (size_t i = 0; i < y.size(); ++i) CHECK(y[i] == y2[i]); + + // error against the float product: each activation is off by at most + // scale/2, so |err| <= scale/2 * sum_k |W| (+ float rounding slack) + for (int t = 0; t < T; ++t) { + const float sa = *reinterpret_cast(&act[(size_t)t * ternary_act_row_bytes(K) + K]); + for (int n = 0; n < N; ++n) { + double ref = 0.0, l1 = 0.0; + for (int k = 0; k < K; ++k) { + const double wv = W[(size_t)n * K + k]; + ref += wv * x[(size_t)t * K + k]; + l1 += std::fabs(wv); + } + const double bound = 0.5 * sa * l1 + 1e-4 * (std::fabs(ref) + 1.0); + CHECK(std::fabs((double)y[(size_t)t * N + n] - ref) <= bound); + } + } + // tight check: scalar reference vs a double product of the dequantized + // weights and the dequantized activations (catches a mis-indexed group scale) + for (int t = 0; t < T; ++t) { + const uint8_t* row = &act[(size_t)t * ternary_act_row_bytes(K)]; + const int8_t* qa = reinterpret_cast(row); + const float sa = *reinterpret_cast(row + K); + for (int n = 0; n < N; ++n) { + double ref = 0.0, mag = 0.0; + for (int k = 0; k < K; ++k) { + const double term = (double)W[(size_t)n * K + k] * ((double)qa[k] * (double)sa); + ref += term; + mag += std::fabs(term); + } + CHECK(std::fabs((double)y[(size_t)t * N + n] - ref) <= 1e-4 * (mag + 1e-6)); + } + } + // every available kernel equals the reference exactly + for (const TernaryKernel* k : ternary_all_kernels()) { + std::vector yk((size_t)T * N, -1.0f); + // one call over all rows, then a split with an uneven boundary + k->fn(w, act.data(), T, yk.data(), 0, N); + for (size_t i = 0; i < y.size(); ++i) + if (y[i] != yk[i]) { std::fprintf(stderr, "kernel %s differs at %zu: %g vs %g (N=%d K=%d T=%d)\n", k->name, i, y[i], yk[i], N, K, T); ++failures; break; } + std::vector ys((size_t)T * N, -1.0f); + const int b = N > 2 ? N / 3 : 1; + k->fn(w, act.data(), T, ys.data(), 0, b); + k->fn(w, act.data(), T, ys.data(), b, N); + for (size_t i = 0; i < y.size(); ++i) CHECK(y[i] == ys[i]); + // the op's split by thread count, including more threads than rows: + // every output written exactly by its own range, nothing else touched + for (int nth : {2, 3, 5, 7, 16, 40}) { + std::vector yt((size_t)T * N, -1.0f); + for (int ith = 0; ith < nth; ++ith) + k->fn(w, act.data(), T, yt.data(), (int)((int64_t)N * ith / nth), (int)((int64_t)N * (ith + 1) / nth)); + for (size_t i = 0; i < y.size(); ++i) CHECK(y[i] == yt[i]); + } + // a single inner range leaves everything outside it untouched + if (N >= 3) { + const int lo = 1, hi = N - 1; + std::vector yr((size_t)T * N, -2.0f); + k->fn(w, act.data(), T, yr.data(), lo, hi); + for (int t = 0; t < T; ++t) + for (int n = 0; n < N; ++n) { + const size_t i = (size_t)t * N + n; + CHECK(n >= lo && n < hi ? yr[i] == y[i] : yr[i] == -2.0f); + } + } + } + // the dispatching entry point equals the reference exactly + std::vector yd((size_t)T * N, -1.0f); + ternary_matmul_rows(w, act.data(), T, yd.data(), 0, N); + for (size_t i = 0; i < y.size(); ++i) CHECK(y[i] == yd[i]); +} + +static void test_silence_is_zero() { + const int N = 4, K = 256, T = 2; + std::vector codes((size_t)N * K, 2); + std::vector s((size_t)N * (K / 128), ggml_fp32_to_fp16(1.0f)); + auto q = pack_trits(codes, N, K); + TernaryWeight w; + ternary_repack(q.data(), s.data(), N, K, w); + std::vector x((size_t)T * K, 0.0f); + std::vector act(ternary_act_row_bytes(K) * T); + ternary_quant_rows(x.data(), K, 0, T, act.data()); + std::vector y((size_t)T * N, 7.0f); + ternary_matmul_rows(w, act.data(), T, y.data(), 0, N); + for (float v : y) CHECK(v == 0.0f && !std::isnan(v)); +} + +// Every quantizer writes the same bytes as ternary_quant_rows_ref. +static void test_quant_identical() { + std::mt19937 rng(99); + std::uniform_real_distribution u(-3.0f, 3.0f); + for (int K : {128, 1024, 4096}) { + const int T = 9; + std::vector x((size_t)T * K); + for (auto& v : x) v = u(rng); + for (int k = 0; k < K; ++k) { + x[(size_t)1 * K + k] = 0.0f; // all zero + x[(size_t)2 * K + k] = (float)(k % 255 - 127) + 0.5f * (k & 1); // amax 127, ties at .5 + x[(size_t)3 * K + k] = u(rng) * 1e-39f; // denormal + x[(size_t)4 * K + k] = u(rng) * 3e38f; // near FLT_MAX + x[(size_t)5 * K + k] = k == 7 ? -0.0f : 0.0f; // negative zero only + } + x[(size_t)2 * K + 5] = 127.0f; + x[(size_t)6 * K + 3] = NAN; // one NaN + x[(size_t)7 * K + 3] = INFINITY; // one inf + x[(size_t)8 * K + K - 1] = -INFINITY; + x[(size_t)8 * K + 0] = NAN; + const size_t rb = ternary_act_row_bytes(K); + std::vector ref(rb * T, 0xAB); + ternary_quant_rows_ref(x.data(), K, 0, T, ref.data()); + for (const TernaryQuant* q : ternary_all_quants()) { + std::vector a(rb * T, 0xCD); + q->fn(x.data(), K, 0, 4, a.data()); // split rows, as the op does + q->fn(x.data(), K, 4, T, a.data()); + for (int t = 0; t < T; ++t) + if (std::memcmp(&a[(size_t)t * rb], &ref[(size_t)t * rb], rb) != 0) { + std::fprintf(stderr, "quant %s differs from the reference (K=%d row %d)\n", q->name, K, t); + ++failures; + } + } + std::vector d(rb * T); + ternary_quant_rows(x.data(), K, 0, T, d.data()); + CHECK(d == ref); + } +} + +// The repack as it was before the lookup table version: an independent scalar +// reference (per byte v % 3, v /= 3, then bit-plane packing one element at a time). +static void repack_reference(const uint8_t* q, const uint16_t* s, int N, int K, TernaryWeight& out) { + const int G = K / kTernaryGroup; + const int nb = (K + 4) / 5; + out.N = N; + out.K = K; + const int B = out.row_blocks(); + out.planes.assign((size_t)B * G * 512, 0); + out.scales.assign((size_t)B * G * kTernaryRowBlock, 0.0f); + std::vector row(K); + for (int n = 0; n < N; ++n) { + const int b = n / kTernaryRowBlock, i = n % kTernaryRowBlock; + for (int g = 0; g < G; ++g) + out.scales[((size_t)b * G + g) * kTernaryRowBlock + i] = ggml_fp16_to_fp32(s[(size_t)n * G + g]); + const uint8_t* qr = q + (size_t)n * nb; + for (int bb = 0; bb < nb; ++bb) { + int v = qr[bb]; + for (int d = 0; d < 5; ++d) { + const int k = bb * 5 + d; + if (k < K) row[k] = (uint8_t)(v % 3); + v /= 3; + } + } + for (int k = 0; k < K; ++k) { + const int g = k / kTernaryGroup, r = k % kTernaryGroup; + const int st = r / 16, p = (r % 16) / 4, j = r % 4; + out.planes[((size_t)b * G + g) * 512 + (size_t)st * 64 + 4 * i + j] |= (uint8_t)(row[k] << (2 * p)); + } + } +} + +static void test_repack_equivalence() { + std::mt19937 rng(77); + int cases = 0; + for (int N : {1, 2, 3, 5, 15, 16, 17, 31, 32, 33, 37, 64, 80, 1024, 4096}) + for (int K : {128, 256, 384, 1024, 4096}) { + if ((size_t)N * K > (size_t)1024 * 1024 && K != 1024) continue; // keep the runtime sane + const int nb = (K + 4) / 5, G = K / kTernaryGroup; + // mode 0 random valid, 1 all 0, 2 all 2, 3 all 1, 4 alternating 0/2, 5 any byte 0..255 (invalid too) + for (int mode = 0; mode < 6; ++mode) { + std::vector q((size_t)N * nb); + for (size_t x = 0; x < q.size(); ++x) { + int v; + switch (mode) { + case 0: v = (int)(rng() % 243); break; + case 1: v = 0; break; + case 2: v = 242; break; // digits 2 2 2 2 2 + case 3: v = 121; break; // digits 1 1 1 1 1 + case 4: v = (x & 1) ? 242 : 0; break; + default: v = (int)(rng() % 256); break; + } + q[x] = (uint8_t)v; + } + std::vector sc((size_t)N * G); + for (size_t x = 0; x < sc.size(); ++x) { + // random F16 bit patterns, including denormals, large values, inf and nan + uint16_t h = (uint16_t)rng(); + if (x % 7 == 0) h = (uint16_t)(rng() % 0x400); // denormal + else if (x % 11 == 0) h = (uint16_t)(0x7000 + rng() % 0xC00); // large + sc[x] = h; + } + TernaryWeight a, b; + repack_reference(q.data(), sc.data(), N, K, a); + ternary_repack(q.data(), sc.data(), N, K, b); + const bool same = a.N == b.N && a.K == b.K && a.planes == b.planes && + a.scales.size() == b.scales.size() && + std::memcmp(a.scales.data(), b.scales.data(), a.scales.size() * 4) == 0; + if (!same) std::fprintf(stderr, "repack mismatch N=%d K=%d mode=%d\n", N, K, mode); + CHECK(same); + ++cases; + } + } + std::printf("repack equivalence: %d cases identical to the reference\n", cases); +} + +int main() { + test_repack_equivalence(); + const Case cases[] = {{8, 128, 1}, {16, 256, 3}, {37, 1024, 5}, {64, 4096, 4}, {5, 384, 7}, {1, 128, 1}, + {48, 1024, 13}, {80, 512, 200}, {40, 1024, 9, true}, {7, 4096, 3, true}}; + unsigned seed = 1; + for (const Case& c : cases) run_case(c, seed++); + // remainder coverage: rows not a multiple of any row tile, columns not a + // multiple of any column tile + for (int N : {1, 2, 3, 5, 7, 37}) + for (int T : {1, 2, 3, 5, 7, 9}) run_case({N, 256, T}, seed++); + test_silence_is_zero(); + test_quant_identical(); + if (failures) return 1; + std::printf("test_ternary: OK (dispatch=%s; tested:", ternary_kernel_name()); + for (const TernaryKernel* k : ternary_all_kernels()) std::printf(" %s", k->name); + std::printf("; quant:"); + for (const TernaryQuant* q : ternary_all_quants()) std::printf(" %s", q->name); + std::printf(")\n"); + return 0; +} diff --git a/tests/test_ternary_load_negative.cpp b/tests/test_ternary_load_negative.cpp new file mode 100644 index 0000000..7859039 --- /dev/null +++ b/tests/test_ternary_load_negative.cpp @@ -0,0 +1,160 @@ +// A doctored packed-ternary GGUF must be refused at load (Model::load returns +// nullptr) while the pristine file loads. Doctored copies are derived from +// $PARAKEET_TEST_GGUF_REDUX_KEEP with the gguf C API and written to the system +// temp directory. Skips (77) when the env var is unset. +#include "model.hpp" +#include "model_loader.hpp" +#include "ternary.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "ggml.h" +#include "gguf.h" + +using namespace pk; + +static int failures = 0; +#define CHECK(c) do { if (!(c)) { std::fprintf(stderr, "FAIL: %s (line %d)\n", #c, __LINE__); ++failures; } } while (0) + +// Removes the file on every exit path. +struct TempFile { + std::string path; + explicit TempFile(const std::string& tag) { + path = (std::filesystem::temp_directory_path() / + ("pk_neg_" + std::to_string((long)getpid()) + "_" + tag + ".gguf")).string(); + } + ~TempFile() { std::remove(path.c_str()); } +}; + +// Rebuilds src through gguf_init_empty (the writer cannot reuse the loaded +// context's tensor infos), skipping tensors for which drop(name) is true, then +// applies edit and writes to tf.path. Returns false on error. +static bool doctor(const char* src, const TempFile& tf, + const std::function& edit, + const std::function& drop = nullptr) { + ggml_context* ctx = nullptr; + gguf_init_params p{/*no_alloc*/ false, &ctx}; + gguf_context* g = gguf_init_from_file(src, p); + if (!g) return false; + gguf_context* w = gguf_init_empty(); + gguf_set_kv(w, g); + const int64_t n = gguf_get_n_tensors(g); + for (int64_t i = 0; i < n; ++i) { + const char* nm = gguf_get_tensor_name(g, i); + if (drop && drop(nm)) continue; + gguf_add_tensor(w, ggml_get_tensor(ctx, nm)); + } + edit(w, ctx); + const bool ok = gguf_write_to_file(w, tf.path.c_str(), false); + gguf_free(w); + gguf_free(g); + ggml_free(ctx); + return ok; +} + +static void expect_refused(const char* src, const std::string& tag, + const std::function& edit) { + TempFile tf(tag); + const bool ok = doctor(src, tf, edit); + CHECK(ok); + if (!ok) return; + auto m = Model::load(tf.path); + std::printf("%s: %s\n", tag.c_str(), m ? "LOADED (unexpected)" : "refused"); + CHECK(m == nullptr); +} + +static bool ends_with(const std::string& s, const char* suf) { + const size_t n = std::strlen(suf); + return s.size() >= n && s.compare(s.size() - n, n, suf) == 0; +} + +int main() { + const char* keep = std::getenv("PARAKEET_TEST_GGUF_REDUX_KEEP"); + if (!keep) { + std::puts("skip: PARAKEET_TEST_GGUF_REDUX_KEEP unset"); + return 77; + } + { + auto m = Model::load(keep); + CHECK(m != nullptr); // pristine file loads + } + // Control: an unedited copy through the same rebuild path must load, so a + // writer regression cannot make every case below pass vacuously. + { + TempFile tf("control"); + CHECK(doctor(keep, tf, [](gguf_context*, ggml_context*) {})); + auto m = Model::load(tf.path); + CHECK(m != nullptr); + ModelLoader ml; + CHECK(ml.load(tf.path)); + CHECK(ternary_flag_consistency_error(ml).empty()); + } + // Flag off with every packed tensor still present. + { + TempFile tf("flag_off"); + CHECK(doctor(keep, tf, [](gguf_context* g, ggml_context*) { + gguf_set_val_bool(g, "parakeet.ternary.present", false); + })); + ModelLoader ml; + CHECK(ml.load(tf.path)); + const std::string err = ternary_flag_consistency_error(ml); + CHECK(err.find("qweight") != std::string::npos); + CHECK(Model::load(tf.path) == nullptr); + } + // Flag off, layer 0 has no packed tensors (dropped), later layers keep + // theirs: the old two-tensor layer-0 probe missed this. + { + TempFile tf("flag_off_later_layer"); + CHECK(doctor(keep, tf, [](gguf_context* g, ggml_context*) { + gguf_set_val_bool(g, "parakeet.ternary.present", false); + }, [](const std::string& n) { + return n.rfind("encoder.layers.0.", 0) == 0 && (ends_with(n, ".qweight") || ends_with(n, ".scales")); + })); + ModelLoader ml; + CHECK(ml.load(tf.path)); + CHECK(!ml.tensor("encoder.layers.0.self_attn.linear_q.qweight")); + CHECK(!ml.tensor("encoder.layers.0.feed_forward1.linear1.qweight")); + CHECK(ml.has_tensor_with_suffix(".qweight")); + const std::string err = ternary_flag_consistency_error(ml); + std::printf("flag_off_later_layer: %s\n", err.c_str()); + CHECK(err.find("qweight") != std::string::npos); + CHECK(Model::load(tf.path) == nullptr); + } + // Flag on but no packed tensors at all. + { + TempFile tf("flag_on_no_packed"); + CHECK(doctor(keep, tf, [](gguf_context*, ggml_context*) {}, + [](const std::string& n) { return ends_with(n, ".qweight") || ends_with(n, ".scales"); })); + ModelLoader ml; + CHECK(ml.load(tf.path)); + CHECK(ternary_flag_consistency_error(ml).find("qweight") != std::string::npos); + CHECK(Model::load(tf.path) == nullptr); + } + expect_refused(keep, "group64", [](gguf_context* g, ggml_context*) { + gguf_set_val_u32(g, "parakeet.ternary.group_size", 64); + }); + expect_refused(keep, "scales_type", [](gguf_context* g, ggml_context*) { + // I8 is half the byte size of F16, so the writer stays inside the buffer. + gguf_set_tensor_type(g, "encoder.layers.0.feed_forward1.linear1.scales", GGML_TYPE_I8); + }); + expect_refused(keep, "ff_dim", [](gguf_context* g, ggml_context*) { + gguf_set_val_u32(g, "parakeet.encoder.ff_dim", 2048); + }); + expect_refused(keep, "vad_nan_frame", [](gguf_context* g, ggml_context*) { + gguf_set_val_f32(g, "parakeet.vad.frame_sec", NAN); + }); + expect_refused(keep, "vad_even_kernel", [](gguf_context* g, ggml_context*) { + gguf_set_val_u32(g, "parakeet.vad.kernel", 4); + }); + if (failures) return 1; + std::puts("test_ternary_load_negative: OK"); + return 0; +} diff --git a/tests/test_ternary_long.cpp b/tests/test_ternary_long.cpp new file mode 100644 index 0000000..beec9cd --- /dev/null +++ b/tests/test_ternary_long.cpp @@ -0,0 +1,119 @@ +// Packed ternary Redux through the local (banded) and chunked attention paths. +// +// Regression for a segfault on long audio: three attention `linear` lambdas +// (build_graph_batched_local, build_graph_local, build_graph_local_chunked) +// read only `.weight`, which does not exist in a packed GGUF, and passed a +// null tensor to ggml_mul_mat. Long audio (encoder frames > 8192, about 655 s) +// switches to local attention automatically; PARAKEET_ATT_CONTEXT forces the same +// path on short audio, which is what this test does, so it stays fast. +// +// Env: PARAKEET_TEST_GGUF_REDUX_KEEP (--ternary keep), PARAKEET_TEST_GGUF_REDUX_DEQ. +// Optional: PARAKEET_TEST_LONG=1 also runs a ~700 s concatenation through the +// automatic trigger (slow, minutes). +#include "model.hpp" +#include "audio_io.hpp" + +#include +#include +#include +#include +#include +#include + +using namespace pk; + +static std::vector words(const std::string& s) { + std::istringstream is(s); + std::vector w; + for (std::string t; is >> t;) w.push_back(t); + return w; +} + +static size_t word_edit_distance(const std::string& a, const std::string& b) { + const auto x = words(a), y = words(b); + std::vector prev(y.size() + 1), cur(y.size() + 1); + for (size_t j = 0; j <= y.size(); ++j) prev[j] = j; + for (size_t i = 1; i <= x.size(); ++i) { + cur[0] = i; + for (size_t j = 1; j <= y.size(); ++j) + cur[j] = std::min({prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x[i - 1] != y[j - 1])}); + std::swap(prev, cur); + } + return prev[y.size()]; +} + +// Within 1 percent of the words, but never less than 2 words: a single split or +// join ("anymore" vs "any more") costs 2 word edits on a short clip. +static bool close_enough(const char* what, const std::string& deq, const std::string& keep) { + const size_t d = word_edit_distance(deq, keep); + const size_t n = words(deq).size(); + const size_t tol = std::max(2, (n + 99) / 100); + std::printf(" %s: word edit distance %zu of %zu words (tolerance %zu)\n", what, d, n, tol); + if (d > tol) { + std::fprintf(stderr, "FAIL: %s differs beyond tolerance\n deq : %s\n keep: %s\n", + what, deq.c_str(), keep.c_str()); + return false; + } + return true; +} + +struct Result { std::string a, b, batch0, batch1; }; + +static bool run(const char* path, const std::vector& sa, const std::vector& sb, + Result& r) { + auto m = Model::load(path); + if (!m) { std::fprintf(stderr, "FAIL: load %s\n", path); return false; } + r.a = m->transcribe_pcm(sa, 16000, Decoder::kTDT); + r.b = m->transcribe_pcm(sb, 16000, Decoder::kTDT); + const auto batch = m->transcribe_pcm_batch({sa, sb}, 16000, Decoder::kTDT); + if (batch.size() != 2) { std::fprintf(stderr, "FAIL: batch size\n"); return false; } + r.batch0 = batch[0]; + r.batch1 = batch[1]; + return true; +} + +int main() { + const char* keep = std::getenv("PARAKEET_TEST_GGUF_REDUX_KEEP"); + const char* deq = std::getenv("PARAKEET_TEST_GGUF_REDUX_DEQ"); + if (!keep || !deq) { + std::puts("skip: PARAKEET_TEST_GGUF_REDUX_KEEP / _DEQ unset"); + return 77; + } + Audio a, b; + if (!load_audio_16k_mono("tests/fixtures/speech.wav", a) || + !load_audio_16k_mono("tests/fixtures/two_speakers.wav", b)) { + std::fprintf(stderr, "FAIL: fixture load\n"); + return 1; + } + + int failures = 0; + // Force local attention (window 64) on short audio: single item goes through + // build_graph_local_chunked, the batch through the batched local path. + setenv("PARAKEET_ATT_CONTEXT", "64", 1); + Result rd, rk; + if (!run(deq, a.samples, b.samples, rd) || !run(keep, a.samples, b.samples, rk)) return 1; + std::printf("local attention, packed: %s\n", rk.a.c_str()); + if (rk.a.empty() || rk.b.empty()) { std::fprintf(stderr, "FAIL: empty transcript\n"); ++failures; } + if (!close_enough("single speech", rd.a, rk.a)) ++failures; + if (!close_enough("single two_speakers", rd.b, rk.b)) ++failures; + if (!close_enough("batched speech", rd.batch0, rk.batch0)) ++failures; + if (!close_enough("batched two_speakers", rd.batch1, rk.batch1)) ++failures; + unsetenv("PARAKEET_ATT_CONTEXT"); + + if (const char* lg = std::getenv("PARAKEET_TEST_LONG"); lg && *lg == '1') { + // Natural trigger: more than 8192 encoder frames (about 655 s). + std::vector longpcm; + while (longpcm.size() < 16000u * 700u) { + longpcm.insert(longpcm.end(), a.samples.begin(), a.samples.end()); + longpcm.insert(longpcm.end(), b.samples.begin(), b.samples.end()); + } + std::string td, tk; + { auto m = Model::load(deq); td = m->transcribe_pcm(longpcm, 16000, Decoder::kTDT); } + { auto m = Model::load(keep); tk = m->transcribe_pcm(longpcm, 16000, Decoder::kTDT); } + if (!close_enough("long 700 s", td, tk)) ++failures; + } + + if (failures) return 1; + std::puts("test_ternary_long: OK"); + return 0; +} diff --git a/tests/test_ternary_model.cpp b/tests/test_ternary_model.cpp new file mode 100644 index 0000000..56ff78d --- /dev/null +++ b/tests/test_ternary_model.cpp @@ -0,0 +1,59 @@ +// Redux with packed ternary weights (native kernel) must decode the reference +// clip to the same transcript as the dequantized Redux GGUF. +// Env: PARAKEET_TEST_GGUF_REDUX_KEEP (--ternary keep), PARAKEET_TEST_GGUF_REDUX_DEQ. +#include "model.hpp" +#include "audio_io.hpp" + +#include +#include +#include + +using namespace pk; + +static const char* kReference = + "Well, I don't wish to see it any more, observed Phoebe, turning away her eyes. " + "It is certainly very like the old portrait."; + +int main() { + const char* keep = std::getenv("PARAKEET_TEST_GGUF_REDUX_KEEP"); + const char* deq = std::getenv("PARAKEET_TEST_GGUF_REDUX_DEQ"); + if (!keep || !deq) { + std::puts("skip: PARAKEET_TEST_GGUF_REDUX_KEEP / _DEQ unset"); + return 77; + } + int failures = 0; + std::string two_deq, two_keep; + for (const char* path : {deq, keep}) { + auto m = Model::load(path); + if (!m) { std::fprintf(stderr, "FAIL: load %s\n", path); return 1; } + const std::string got = m->transcribe_path("tests/fixtures/speech.wav", Decoder::kTDT); + std::printf("%s\n -> %s\n", path, got.c_str()); + if (got != kReference) { std::fprintf(stderr, "FAIL: transcript differs for %s\n", path); ++failures; } + const std::string two = m->transcribe_path("tests/fixtures/two_speakers.wav", Decoder::kTDT); + std::printf(" two_speakers -> %s\n", two.c_str()); + (path == keep ? two_keep : two_deq) = two; + if (path == keep) { + // batched == per-item on the packed model + Audio a, b; + if (!load_audio_16k_mono("tests/fixtures/speech.wav", a) || + !load_audio_16k_mono("tests/fixtures/two_speakers.wav", b)) { + std::fprintf(stderr, "FAIL: fixture load\n"); + return 1; + } + const auto batch = m->transcribe_pcm_batch({a.samples, b.samples}, 16000, Decoder::kTDT); + if (batch.size() != 2 || batch[0] != got || batch[1] != two) { + std::fprintf(stderr, "FAIL: batched transcripts differ from per-item\n"); + for (const auto& t : batch) std::fprintf(stderr, " batch: %s\n", t.c_str()); + ++failures; + } + } + } + if (two_deq != two_keep) { + std::fprintf(stderr, "FAIL: two_speakers transcript differs between dequantized and packed Redux\n" + " deq : %s\n keep: %s\n", two_deq.c_str(), two_keep.c_str()); + ++failures; + } + if (failures) return 1; + std::puts("test_ternary_model: OK"); + return 0; +} diff --git a/tests/test_transcribe_vad.cpp b/tests/test_transcribe_vad.cpp new file mode 100644 index 0000000..a3e324c --- /dev/null +++ b/tests/test_transcribe_vad.cpp @@ -0,0 +1,160 @@ +// VAD-segmented transcription: short audio unchanged, long audio close to the +// single-pass transcript, timestamps monotonic, C-API error path. +// Env: PARAKEET_TEST_GGUF_ULTRA (has a VAD head), PARAKEET_TEST_GGUF (does not; optional). +#include +#include +#include +#include +#include +#include +#include +#include + +#include "audio_io.hpp" +#include "model.hpp" +#include "parakeet_capi.h" + +using namespace pk; + +static int failures = 0; +#define CHECK(c) do { if (!(c)) { std::fprintf(stderr, "FAIL: %s (line %d)\n", #c, __LINE__); ++failures; } } while (0) + +static std::vector words_of(const std::string& s) { + std::istringstream is(s); + std::vector w; + for (std::string x; is >> x;) { + std::string y; + for (char c : x) if (std::isalnum((unsigned char)c)) y += (char)std::tolower((unsigned char)c); + if (!y.empty()) w.push_back(y); + } + return w; +} +static size_t edit_distance(const std::vector& a, const std::vector& b) { + std::vector prev(b.size() + 1), cur(b.size() + 1); + for (size_t j = 0; j <= b.size(); ++j) prev[j] = j; + for (size_t i = 1; i <= a.size(); ++i) { + cur[0] = i; + for (size_t j = 1; j <= b.size(); ++j) + cur[j] = std::min({prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (a[i - 1] == b[j - 1] ? 0 : 1)}); + std::swap(prev, cur); + } + return prev[b.size()]; +} + +int main() { + const char* ultra = std::getenv("PARAKEET_TEST_GGUF_ULTRA"); + const char* plain = std::getenv("PARAKEET_TEST_GGUF"); + if (!ultra) { std::puts("skip: PARAKEET_TEST_GGUF_ULTRA unset"); return 77; } + + auto m = Model::load(ultra); + if (!m) { std::fprintf(stderr, "FAIL: load\n"); return 1; } + + // 1. short audio: identical to the plain path + Audio shortclip; + CHECK(load_audio_16k_mono("tests/fixtures/speech.wav", shortclip)); + const std::string a = m->transcribe_pcm(shortclip.samples, 16000, Decoder::kTDT); + const std::string b = m->transcribe_pcm_vad(shortclip.samples, 16000, Decoder::kTDT); + CHECK(a == b); + + // 2. self-contained long clip (~62 s) from tracked fixtures + Audio clip; + { + const char* parts[] = {"tests/fixtures/two_speakers.wav", "tests/fixtures/speech.wav", + "tests/fixtures/two_speakers.wav", "tests/fixtures/speech.wav"}; + for (const char* f : parts) { + Audio x; + CHECK(load_audio_16k_mono(f, x)); + clip.samples.insert(clip.samples.end(), x.samples.begin(), x.samples.end()); + } + clip.sample_rate = 16000; + } + const double total = (double)clip.samples.size() / 16000.0; + { + SegmenterOpts o; + o.frame_sec = m->config().vad.frame_sec; + const auto segs = segment_by_vad(m->vad_probabilities(clip.samples), total, o); + std::printf("62 s clip (%.2f s): %zu segments:", total, segs.size()); + for (const auto& sg : segs) std::printf(" [%.2f-%.2f]", sg.start, sg.end); + std::printf("\n"); + CHECK(segs.size() >= 2); + if (!segs.empty()) { + CHECK(segs.front().start == 0.0); + CHECK(std::fabs(segs.back().end - total) < 1e-6); + for (size_t i = 0; i < segs.size(); ++i) { + CHECK(segs[i].end - segs[i].start <= o.max_seg_sec + 1e-6); + if (i + 1 < segs.size()) CHECK(segs[i].end == segs[i + 1].start); + } + } + } + { + const std::string full = m->transcribe_pcm(clip.samples, 16000, Decoder::kTDT); + const std::string seg = m->transcribe_pcm_vad(clip.samples, 16000, Decoder::kTDT); + const auto wf = words_of(full), ws = words_of(seg); + const double diff = (double)edit_distance(wf, ws) / (double)std::max(1, wf.size()); + std::printf("62 s clip: %zu vs %zu words, word diff %.3f\n", wf.size(), ws.size(), diff); + CHECK(!ws.empty()); + CHECK(diff < 0.15); + const Transcription tr = m->transcribe_pcm_vad_with_timestamps(clip.samples, 16000, Decoder::kTDT); + CHECK(!tr.words.empty()); + float prev = -1.0f; + for (const Word& w : tr.words) { + CHECK(w.start >= prev - 1e-3f); + CHECK(w.end >= w.start); + prev = w.start; + } + if (!tr.words.empty()) CHECK(tr.words.back().end <= (float)total + 0.5f); + for (size_t i = 1; i < tr.tokens.size(); ++i) CHECK(tr.tokens[i].frame >= tr.tokens[i - 1].frame); + CHECK(words_of(tr.text) == ws); + } + + // 3. optional 180 s clip (untracked file) + Audio longclip; + if (!load_audio_16k_mono("benchmarks/audio/diverse/i_have_a_dream.wav", longclip)) { + std::puts("skip long clip: file missing"); + } else { + const std::string full = m->transcribe_pcm(longclip.samples, 16000, Decoder::kTDT); + const std::string seg = m->transcribe_pcm_vad(longclip.samples, 16000, Decoder::kTDT); + const auto wf = words_of(full), ws = words_of(seg); + const double diff = (double)edit_distance(wf, ws) / (double)std::max(1, wf.size()); + std::printf("long clip: %zu vs %zu words, word diff %.3f\n", wf.size(), ws.size(), diff); + CHECK(!ws.empty()); + CHECK(diff < 0.08); + const Transcription tr = m->transcribe_pcm_vad_with_timestamps(longclip.samples, 16000, Decoder::kTDT); + const float dur = (float)longclip.samples.size() / 16000.0f; + CHECK(!tr.words.empty()); + float prev = -1.0f; + for (const Word& w : tr.words) { + CHECK(w.start >= prev - 1e-3f); + CHECK(w.end >= w.start); + CHECK(w.end <= dur + 0.5f); + prev = w.start; + } + CHECK(words_of(tr.text) == ws); + } + + // 4. C-API: JSON works on a VAD model, NULL + message on a model without one + parakeet_ctx* ctx = parakeet_capi_load(ultra); + CHECK(ctx != nullptr); + if (ctx) { + char* j = parakeet_capi_transcribe_path_json_vad(ctx, "tests/fixtures/speech.wav", 2); + CHECK(j != nullptr); + if (j) { CHECK(std::string(j).find("\"words\"") != std::string::npos); parakeet_capi_free_string(j); } + parakeet_capi_free(ctx); + } + if (!plain) std::puts("note: PARAKEET_TEST_GGUF unset, skipping no-VAD-head error check"); + if (plain) { + auto pm = Model::load(plain); + CHECK(pm && !pm->config().vad.present); + parakeet_ctx* c2 = parakeet_capi_load(plain); + CHECK(c2 != nullptr); + if (c2) { + char* j = parakeet_capi_transcribe_path_json_vad(c2, "tests/fixtures/speech.wav", 2); + CHECK(j == nullptr); + CHECK(std::string(parakeet_capi_last_error(c2)).find("VAD") != std::string::npos); + parakeet_capi_free(c2); + } + } + if (failures) return 1; + std::puts("test_transcribe_vad: OK"); + return 0; +} diff --git a/tests/test_vad_head.cpp b/tests/test_vad_head.cpp new file mode 100644 index 0000000..7f5641b --- /dev/null +++ b/tests/test_vad_head.cpp @@ -0,0 +1,127 @@ +// VadHead math on hand-checkable weights, plus a smoke test on a real model. +#include "vad_head.hpp" + +#include +#include +#include +#include + +#include "model_loader.hpp" + +using namespace pk; + +static int failures = 0; +#define CHECK(c) do { if (!(c)) { std::fprintf(stderr, "FAIL: %s (line %d)\n", #c, __LINE__); ++failures; } } while (0) +static bool near(float a, float b) { return std::fabs(a - b) < 1e-5f; } +static float sig(float z) { return 1.0f / (1.0f + std::exp(-z)); } + +// d_in = 2, hidden = 2, kernel = 3, deliberately asymmetric so a transposed +// [in][out] indexing bug cannot pass. +// proj: h0 = x0 + 2*x1 + 0.5, h1 = x1 - 0.5 +// ctx (taps kk = 0, 1, 2 are t-1, t, t+1): +// o0 = 1 * in1(t) + 2 * in0(t+1) + 0.25 (o != i, and a right-hand tap) +// o1 = 1 * in0(t-1) - 0.5 (o != i, left tap) +// out: z = 0.1 * o0 + 0.05 * o1 - 0.5 +static VadWeights tiny() { + VadWeights w; + w.d_in = 2; w.hidden = 2; w.kernel = 3; + w.proj_w = {1, 2, 0, 1}; + w.proj_b = {0.5f, -0.5f}; + w.ctx_w.assign(2 * 2 * 3, 0.0f); + w.ctx_w[(0 * 2 + 1) * 3 + 1] = 1.0f; // o=0, i=1, kk=1 (center) + w.ctx_w[(0 * 2 + 0) * 3 + 2] = 2.0f; // o=0, i=0, kk=2 (t+1) + w.ctx_w[(1 * 2 + 0) * 3 + 0] = 1.0f; // o=1, i=0, kk=0 (t-1) + w.ctx_b = {0.25f, -0.5f}; + w.out_w = {0.1f, 0.05f}; + w.out_b = -0.5f; + return w; +} + +// x = (1,0), (0,1), (2,1). After proj, before ReLU: +// t0: (1.5, -0.5) t1: (2.5, 0.5) t2: (4.5, 0.5) +// after ReLU h1: t0 (1.5, 0), t1 (2.5, 0.5), t2 (4.5, 0.5). +// ctx pre-activation (right pad and left pad are zero): +// t0: o0 = h1[0][1] + 2*h1[1][0] + 0.25 = 0 + 5 + 0.25 = 5.25 +// o1 = (left pad 0) - 0.5 = -0.5 +// t1: o0 = 0.5 + 2*4.5 + 0.25 = 9.75 o1 = h1[0][0] - 0.5 = 1.0 +// t2: o0 = 0.5 + (right pad 0) + 0.25 = 0.75 o1 = h1[1][0] - 0.5 = 2.0 +static const float kX[] = {1, 0, 0, 1, 2, 1}; + +static void test_plain_variant() { + VadVariant v; // relu after proj, no residual, relu after ctx + v.relu_after_proj = true; v.residual = false; v.relu_after_ctx = true; + const auto p = VadHead::run(tiny(), v, kX, 3); + CHECK(p.size() == 3); + // ReLU after ctx: t0 (5.25, 0), t1 (9.75, 1.0), t2 (0.75, 2.0) + // z0 = 0.1*5.25 + 0.05*0 - 0.5 = 0.025 + // z1 = 0.1*9.75 + 0.05*1.0 - 0.5 = 0.525 + // z2 = 0.1*0.75 + 0.05*2.0 - 0.5 = -0.325 + CHECK(near(p[0], sig(0.025f))); + CHECK(near(p[1], sig(0.525f))); + CHECK(near(p[2], sig(-0.325f))); +} + +static void test_residual_variant() { + VadVariant v; + v.relu_after_proj = true; v.residual = true; v.relu_after_ctx = true; + const auto p = VadHead::run(tiny(), v, kX, 3); + // Residual adds h1 (post-ReLU) before the ReLU after ctx: + // t0: (5.25+1.5, -0.5+0) = (6.75, -0.5 -> 0) + // t1: (9.75+2.5, 1.0+0.5) = (12.25, 1.5) + // t2: (0.75+4.5, 2.0+0.5) = (5.25, 2.5) + // z0 = 0.675 - 0.5 = 0.175 + // z1 = 1.225 + 0.075 - 0.5 = 0.8 + // z2 = 0.525 + 0.125 - 0.5 = 0.15 + CHECK(near(p[0], sig(0.175f))); + CHECK(near(p[1], sig(0.8f))); + CHECK(near(p[2], sig(0.15f))); +} + +static void test_relu_matters() { + const float x[] = {-1, -2}; // T = 1, negative inputs + VadVariant with, without; + with.relu_after_proj = true; with.relu_after_ctx = true; + without.relu_after_proj = false; without.relu_after_ctx = false; + const auto a = VadHead::run(tiny(), with, x, 1); + const auto b = VadHead::run(tiny(), without, x, 1); + // with ReLUs: proj pre (-1-4+0.5, -2-0.5) = (-4.5, -2.5) -> h1 = (0, 0); + // ctx (0.25, -0.5) -> ReLU (0.25, 0); z = 0.025 - 0.5 = -0.475 + CHECK(near(a[0], sig(-0.475f))); + // without: h1 = (-4.5, -2.5). T = 1 so only the center tap is in range: + // o0 = in1 + 0.25 = -2.5 + 0.25 = -2.25, o1 = -0.5 (its tap is padding) + // z = 0.1*-2.25 + 0.05*-0.5 - 0.5 = -0.225 - 0.025 - 0.5 = -0.75 + CHECK(near(b[0], sig(-0.75f))); +} + +static void test_from_index() { + const VadVariant v = VadVariant::from_index(5); + CHECK(v.relu_after_proj && !v.residual && v.relu_after_ctx); + const VadVariant z = VadVariant::from_index(0); + CHECK(!z.relu_after_proj && !z.residual && !z.relu_after_ctx); +} + +static void test_real_model() { + const char* p = std::getenv("PARAKEET_TEST_GGUF_ULTRA"); + if (!p) return; + ModelLoader ml; + CHECK(ml.load(p)); + VadHead head(ml); + CHECK(std::fabs(head.frame_sec() - 0.08f) < 1e-6f); + const int T = 40, D = (int)ml.config().vad.d_in; + std::vector x((size_t)T * D); + for (size_t i = 0; i < x.size(); ++i) x[i] = std::sin(0.01f * (float)i); + const auto prob = head.probabilities(x.data(), T); + CHECK((int)prob.size() == T); + for (float v : prob) CHECK(std::isfinite(v) && v >= 0.0f && v <= 1.0f); +} + +int main() { + test_plain_variant(); + test_residual_variant(); + test_relu_matters(); + test_from_index(); + test_real_model(); + if (failures) return 1; + std::puts("test_vad_head: OK"); + return 0; +} diff --git a/tests/test_vad_segmenter.cpp b/tests/test_vad_segmenter.cpp new file mode 100644 index 0000000..cd96a7b --- /dev/null +++ b/tests/test_vad_segmenter.cpp @@ -0,0 +1,238 @@ +#include "vad_segmenter.hpp" + +#include +#include +#include +#include +#include +#include + +using namespace pk; + +static int failures = 0; +#define CHECK(c) do { if (!(c)) { std::fprintf(stderr, "FAIL: %s (line %d)\n", #c, __LINE__); ++failures; } } while (0) + +static const double FS = 0.08; + +// n frames of speech (p = 0.95) with silent ranges [a, b) in frames (p = 0.02). +static std::vector make_p(int n, std::initializer_list> silences) { + std::vector p((size_t)n, 0.95f); + for (auto s : silences) + for (int i = s.first; i < s.second && i < n; ++i) p[(size_t)i] = 0.02f; + return p; +} + +static void check_tiling(const std::vector& s, double total, double max_seg) { + CHECK(!s.empty()); + if (s.empty()) return; + CHECK(s.front().start == 0.0); + CHECK(std::fabs(s.back().end - total) < 1e-9); + for (size_t i = 0; i < s.size(); ++i) { + CHECK(s[i].end > s[i].start); + CHECK(s[i].end - s[i].start <= max_seg + 1e-9); + if (i + 1 < s.size()) CHECK(s[i].end == s[i + 1].start); + } +} + +static void test_short_is_single() { + SegmenterOpts o; + auto s = segment_by_vad(make_p(250, {}), 20.0, o); + CHECK(s.size() == 1 && s[0].start == 0.0 && s[0].end == 20.0); + s = segment_by_vad({}, 0.5, o); // shorter than one frame + CHECK(s.size() == 1 && s[0].end == 0.5); + s = segment_by_vad(make_p(375, {}), 30.0, o); // exactly the cap + CHECK(s.size() == 1 && s[0].end == 30.0); +} + +static void test_all_speech_hard_cuts() { + SegmenterOpts o; + const double total = 70.0; + auto s = segment_by_vad(make_p((int)std::ceil(total / FS), {}), total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 3); +} + +static void test_cuts_land_in_pauses() { + SegmenterOpts o; + const double total = 70.0; + // pauses at [312,325) and [640,653): midpoints 318 (25.44 s) and 646 (51.68 s) + auto p = make_p((int)std::ceil(total / FS), {{312, 325}, {640, 653}}); + auto s = segment_by_vad(p, total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 3); + if (s.size() == 3) { + CHECK(std::fabs(s[0].end - 25.44) < 0.1); + CHECK(std::fabs(s[1].end - 51.68) < 0.1); + } +} + +static void test_picks_longest_pause() { + SegmenterOpts o; + const double total = 40.0; + // two pauses in the last third of the first window: 22-22.4 s (short) and 27-28 s (long) + auto p = make_p((int)std::ceil(total / FS), {{275, 280}, {337, 350}}); + auto s = segment_by_vad(p, total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 2); + if (s.size() == 2) CHECK(std::fabs(s[0].end - 27.5) < 0.3); +} + +static void test_all_silence() { + SegmenterOpts o; + const double total = 65.0; + auto s = segment_by_vad(make_p((int)std::ceil(total / FS), {{0, 100000}}), total, o); + check_tiling(s, total, o.max_seg_sec); +} + +static void test_pause_shorter_than_min_is_ignored() { + SegmenterOpts o; + const double total = 45.0; + // a 2-frame (0.16 s) pause at 25 s is below min_pause_sec (0.32 s): hard cut at 30 s + auto s = segment_by_vad(make_p((int)std::ceil(total / FS), {{312, 314}}), total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 2 && std::fabs(s[0].end - 30.0) < 0.1); +} + +static void test_random_property() { + std::mt19937 rng(7); + std::uniform_real_distribution u(0.0f, 1.0f); + std::uniform_int_distribution run_len(1, 12); + SegmenterOpts o; + for (int trial = 0; trial < 200; ++trial) { + const double total = 31.0 + 600.0 * u(rng); + const int n = (int)std::ceil(total / FS); + std::vector p((size_t)n, 0.9f); + // Generate bursty silence: runs of 1 to 12 frames + for (int i = 0; i < n; ) { + if (u(rng) < 0.2f) { // start a silence run + int len = run_len(rng); + for (int j = 0; j < len && i + j < n; ++j) p[(size_t)(i + j)] = 0.05f; + i += len; + } else { + ++i; + } + } + check_tiling(segment_by_vad(p, total, o), total, o.max_seg_sec); + } +} + +static void test_fallback_window() { + // Pause at frames [150,163) only (in first window but not in last third). + // First search [250,375) finds nothing; fallback search [100,375) finds [150,163). + // Midpoint 156 => 12.48 s. Total 42.4 s (531 frames) yields exactly 2 segments. + SegmenterOpts o; + const double total = 42.4; + auto p = make_p((int)std::ceil(total / FS), {{150, 163}}); + auto s = segment_by_vad(p, total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 2); + if (s.size() == 2) CHECK(std::fabs(s[0].end - 12.48) < 0.1); +} + +static void test_min_seg_respected() { + // Pause at frames [50,63) (4.0 to 5.04 s) only. + // Fallback window starts at min_seg frame 100, so pause is ignored. + // Expected hard cut at 30.0 s. + SegmenterOpts o; + const double total = 45.0; + auto p = make_p((int)std::ceil(total / FS), {{50, 63}}); + auto s = segment_by_vad(p, total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 2 && std::fabs(s[0].end - 30.0) < 0.1); +} + +static void test_tie_goes_to_later_run() { + // Two equal 13-frame pauses at [270,283) and [330,343), both in last third. + // When tied, later run wins. Midpoint of [330,343) is 336 => 26.88 s. + SegmenterOpts o; + const double total = 45.0; + auto p = make_p((int)std::ceil(total / FS), {{270, 283}, {330, 343}}); + auto s = segment_by_vad(p, total, o); + check_tiling(s, total, o.max_seg_sec); + CHECK(s.size() == 2); + if (s.size() == 2) CHECK(std::fabs(s[0].end - 26.88) < 0.1); +} + +static void test_exact_totals() { + // Test 1: total 60.0 s (no pauses, all speech). + // Expected: exactly two segments [0, 30] and [30, 60]. + SegmenterOpts o; + auto s = segment_by_vad(make_p(750, {}), 60.0, o); + check_tiling(s, 60.0, o.max_seg_sec); + CHECK(s.size() == 2); + if (s.size() == 2) { + CHECK(std::fabs(s[0].end - 30.0) < 1e-6); + CHECK(std::fabs(s[1].end - 60.0) < 1e-6); + } + + // Test 2: total 30.08 s with 376 frames (no pauses). + // Expected: two segments with s[0].end at 30.0 and last segment about 0.08 s. + s = segment_by_vad(make_p(376, {}), 30.08, o); + check_tiling(s, 30.08, o.max_seg_sec); + CHECK(s.size() == 2); + if (s.size() == 2) { + CHECK(std::fabs(s[0].end - 30.0) < 1e-6); + CHECK(std::fabs((s[1].end - s[1].start) - 0.08) < 1e-6); + } +} + +static void test_degenerate_opts() { + const double bad[] = {0.0, -0.08, INFINITY, NAN}; + for (double fs : bad) { + SegmenterOpts o; + o.frame_sec = fs; + auto s = segment_by_vad(std::vector(100, 0.0f), 100.0, o); + CHECK(s.size() == 1); + if (s.size() == 1) CHECK(s[0].start == 0.0 && s[0].end == 100.0); + } + SegmenterOpts o; + o.max_seg_sec = 0.05; // smaller than a frame + auto s = segment_by_vad(std::vector(100, 0.0f), 100.0, o); + CHECK(s.size() == 1); + o = SegmenterOpts(); + o.max_seg_sec = NAN; + s = segment_by_vad(std::vector(100, 0.0f), 100.0, o); + CHECK(s.size() == 1); +} + +static void test_nonfinite_and_huge_opts() { + const double bads[] = {INFINITY, -INFINITY, NAN, 1e300, 1e7}; + for (double v : bads) { + for (int which = 0; which < 3; ++which) { + SegmenterOpts o; + if (which == 0) o.min_seg_sec = v; + else if (which == 1) o.min_pause_sec = v; + else o.max_seg_sec = v; + auto s = segment_by_vad(std::vector(2000, 0.0f), 100.0, o); + CHECK(s.size() == 1); + if (s.size() == 1) CHECK(s[0].start == 0.0 && s[0].end == 100.0); + } + } + const float thr[] = {INFINITY, -INFINITY, NAN}; + for (float t : thr) { + SegmenterOpts o; + o.threshold = t; + auto s = segment_by_vad(std::vector(2000, 0.0f), 100.0, o); + CHECK(s.size() == 1); + if (s.size() == 1) CHECK(s[0].start == 0.0 && s[0].end == 100.0); + } +} + +int main() { + test_nonfinite_and_huge_opts(); + test_degenerate_opts(); + test_short_is_single(); + test_all_speech_hard_cuts(); + test_cuts_land_in_pauses(); + test_picks_longest_pause(); + test_all_silence(); + test_pause_shorter_than_min_is_ignored(); + test_fallback_window(); + test_min_seg_respected(); + test_tie_goes_to_later_run(); + test_exact_totals(); + test_random_property(); + if (failures) return 1; + std::puts("test_vad_segmenter: OK"); + return 0; +}