parakeet-mlx: bf16 precision, native load, distributable bf16 model, download progress + dedup#8
Open
fiorelorenzo wants to merge 15 commits into
Open
Conversation
Adds `Model::transcribe_stream()` returning a `Box<dyn StreamSession>` that lets callers push audio incrementally and receive partial transcripts. v1 implements re-encode-growing-buffer for the Parakeet ONNX engine: each push recomputes mel + runs the encoder on the full buffer-so-far, then resumes the greedy TDT decode from the previous time-step. Quadratic in audio length; acceptable for short utterances (typical dictation <30s on M-series M4 encodes in tens of ms per pass), to be optimized in a follow-up via chunked-attention encoder state. `Engine` trait grows a `transcribe_stream` method with a default `Err(NotSupported)` impl, so Whisper / qwen3-asr / parakeet-mlx engines compile unchanged. Streaming support for those is a follow-up. Added `tests/parakeet_streaming.rs` (ignored by default, needs cached weights) asserting the chunked transcript matches the one-shot transcript on the same audio. Consumer: the Lirevo dictation app (https://github.com/fiorelorenzo/lirevo) needs this for its live-overlay UX.
The v1 streaming API committed tokens off whatever encoder context was available at the time of each push, and never revisited those tokens when later pushes grew the encoder's input. Since Parakeet's Conformer encoder is non-causal, the encoder output for the same early frames is not stable across pushes, and on utterances longer than ~20s the cumulative streamed text could drift from the one-shot transcribe() result on the same audio. v1.1 always re-decodes from t=0 every push and computes the consumer- facing `delta` as the diff vs the previously-emitted text. The cumulative text in `PartialTranscript.text` is now byte-identical to what `transcribe()` would return on the same audio at every push. The UI should treat `text` as authoritative and use `delta` as a hint; early tokens may be corrected retroactively when later pushes give the encoder more context. Removed the per-session TDT state carry-across (`state_h`, `state_c`, `last_token`, `emitted`) - no longer needed since each push decodes from scratch. Updated the integration test to assert strict text equality between streamed cumulative text and one-shot text. Build, clippy, and fmt gates green.
Mirrors the Parakeet ONNX streaming impl from the previous commits on this branch onto the MLX engine. Same re-decode-growing-buffer architecture, same PartialTranscript contract, same byte-equivalence to one-shot `transcribe()` by construction. Refactors ParakeetMlxEngine to hold `conformer`, `predict`, `joint` behind `Arc<Mutex<_>>` so the streaming session can hold cloned handles, paralleling the ONNX engine's `Arc<Mutex<ort::Session>>` pattern. Vocab, durations, mel_config and time_ratio are cloned into the session by value. The mel + encoder + TDT decode pipeline is pulled out of the engine's `transcribe` into `run_full_pipeline` so the streaming session can reuse the same code path. `mlx_synchronize` is now called via a small `gpu_synchronize` helper at the end of `transcribe`, `push`, and `finish` to keep the GPU-buffer-lifetime contract intact. Send on `ParakeetMlxStreamSession` is asserted unsafe with the same single-thread-at-a-time invariant the engine documents: the trait surface holds the session behind `&mut self` and `Box<dyn StreamSession + Send>`. Extends `tests/parakeet_streaming.rs` with a parakeet-mlx variant gated on the `parakeet-mlx` feature. The file's outer cfg switched from `feature = "parakeet"` to `any(parakeet, parakeet-mlx)` so each engine's test can be enabled independently.
rc.10 pinned smallvec to a 2.0.0-alpha.10 prerelease; rc.12 uses smallvec ^1.15 which resolves a downstream resolver collision for projects that pull mistral.rs (whose transitive serde-saphyr requires smallvec ^2.0.0-alpha.12, a different prerelease that is mutually exclusive with rc.10's pin). rc.12 also moved to ndarray 0.17 (audiopipe was on 0.16) and made `ort::Error` generic on a recovery payload — added a `From<ort::Error<SessionBuilder>>` bridge in `error.rs` so `?` keeps working at session-builder call sites. No behaviour change in audiopipe's parakeet / whisper paths; build is clean against `parakeet,parakeet-mlx,qwen3-asr-ggml,whisper,metal,coreml`.
Streaming re-decodes the whole (growing) audio buffer on every push, and MLX caches Metal buffers keyed by size. Dictations of increasing length keep minting fresh buffer sizes the cache never reuses, so the MLX cache and the process IOAccelerator footprint climb unbounded across a session (~11 GB after a dozen dictations observed in the Lirevo host). Add the MLX memory FFI (clear_cache / set_cache_limit / get_*_memory) and: - set a 512 MiB cache ceiling once at engine init (AUDIOPIPE_MLX_CACHE_LIMIT_MB to tune, =0 to disable); - clear_cache after finish() and batch transcribe() (AUDIOPIPE_NO_MLX_CACHE_CLEAR to opt out). Validated with examples/mlx_stream_leak_probe over increasing-length dictations: mlx_cache 58 -> 968 -> ... -> 2959 MB (unbounded) before, flat at 0 MB after; phys_footprint 6.0 GB -> 3.0 GB (just the model).
…right away Without the eval, each bf16 weight stayed a lazy cast node holding a reference to its fp32 input, so the full ~2.4 GB fp32 copy of the weights lingered resident from load until the first inference walked the graph. A just-loaded, idle model therefore sat at the fp32 footprint instead of bf16. Forcing the cast at load lets the fp32 source drop with the safetensors map, so the idle model settles at the bf16 footprint immediately. Probe (parakeet-tdt-0.6b-v3-mlx): baseline mlx_active 2430 -> 1235 MB; steady state unchanged at 1235 MB.
Author
|
Added ae35e13: the load-time Without it the bf16 result of each weight cast stayed a lazy node holding a Probe (parakeet-tdt-0.6b-v3-mlx): baseline |
…nd-trip load_safetensors now builds the mlx Array directly in its source dtype for BF16 and F16 inputs, using the native half::bf16 / half::f16 ArrayElement impls in mlx-rs 0.25. Previously both branches up-converted to a Vec<f32> and built an f32 Array, so loading a bf16 model materialized a full fp32 copy of the weights before to_weight_dtype cast it back down. The f32 branch is unchanged; the U32 skip is unchanged.
…tribution) Offline, pure-Rust converter (no python, no mlx) that reads <input_dir>/model.safetensors and writes a bf16 copy to <output_dir>: F32 and F16 tensors are down-converted to BF16, BF16 passes through, any other dtype is copied verbatim, and __metadata__ plus config.json / vocab.txt / tokenizer.model are carried over so the output dir is a complete loadable model. Halves the on-disk model (2.34 GiB -> 1.17 GiB for parakeet-tdt-0.6b-v3). An optional --verify flag loads the result via ParakeetMlxEngine::from_dir and reports mlx active memory.
…check example Model::from_dir can now load a local Parakeet-MLX directory (e.g. a pre-converted bf16 model), routing to ParakeetMlxEngine. The parity_check example transcribes the same audio via the HF-cache fp32 model and a local bf16 dir and asserts identical text - used to validate convert_to_bf16 output. Verified on jfk.wav: byte-identical transcript (bf16 weights are the same whether cast at runtime or pre-stored).
…ET_MLX_REPO Lets a consumer point the Parakeet-MLX loader at a drop-in alternative repo (e.g. a pre-converted bf16 repo for half the download) without forking the crate. Defaults to the upstream mlx-community fp32 repo. Validated: probe loads the bf16 repo end-to-end with the env var set (mlx_active 1235 MB, no panic).
Repeated download triggers (warm-up + reloads on a fresh cache) each spawned a thread that re-downloaded the same files. Concurrent downloads fight over hf_hub's per-blob lock: all but the lock holder fail every retry and log a spurious 'download/load failed', even though the holder completes. Track which model names have a download in flight and skip duplicate spawns, so a fresh install downloads each model exactly once with no lock contention.
… reporting Add ParakeetMlxEngine::download_with_progress and a ProgressAdapter that bridges hf-hub's Progress trait onto an FnMut(received, total) closure, plus Model::download_pretrained_with_progress routing by name. Lets a consumer (the Tauri onboarding wizard) show a real byte progress bar while the STT model's model.safetensors downloads; small metadata files are fetched without progress.
fiorelorenzo
changed the title
Contributor
|
@fiorelorenzo please benchmark bf16 with fp32 in my evals it reduce a lot accuracy |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This branch collects the Parakeet-MLX work my downstream app (Lirevo) needs. It
is stacked on #7 (streaming), so GitHub's diff against
mainalso shows #7'scommits plus a Metal-cache-bound fix and an
ortbump that rode on that branch.Happy to rebase / split for upstream once #7 lands — see "Splitting" below.
What's here (parakeet-mlx)
Memory / precision:
bb58823load weights in bf16 (halves resident memory)549c2a2down-cast attention weights + make forward reads bf16-tolerantee230f5ParakeetPrecisionload param (Int8 reserved -> falls back to bf16)ae35e13eval bf16 weights at load so the fp32 copy frees immediately insteadof lingering until the first inference (baseline mlx_active 2430 -> 1235 MB)
d7004d6load bf16/f16 safetensors natively, no f32 round-tripDistribution:
944408aconvert_to_bf16example (fp32 safetensors -> bf16, half the size)940d983parakeet-mlxarm onModel::from_dir+parity_checkexample3538b17AUDIOPIPE_PARAKEET_MLX_REPOenv override (point at a pre-convertedbf16 repo for half the download, default unchanged)
First-run download robustness:
6869589dedup in-flightspawn_pretrained_download(concurrent triggers nolonger fight over hf_hub's per-blob lock and log spurious failures)
5fb8a00download_with_progressso a consumer can show a real progress barValidation
(verified on real speech via
parity_checkon jfk.wav).Splitting
If you'd prefer smaller PRs, the natural cut is: (1) bf16 precision
(
bb58823/549c2a2/ee230f5/ae35e13/d7004d6), (2) distribution(
944408a/940d983/3538b17), (3) download robustness(
6869589/5fb8a00). I can rebase each ontomainafter #7 merges.Disclosure
Parts of this branch were written with AI assistance. I reviewed the diffs, ran
the probe + parity checks, and run it through a real app.